/
embedding.js
67 lines (61 loc) · 1.49 KB
/
embedding.js
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
import Matrix from '../../../util/matrix.js'
import Tensor from '../../../util/tensor.js'
import Layer from './base.js'
/**
* Embedding layer
*/
export default class EmbeddingLayer extends Layer {
/**
* @param {object} config object
* @param {number} [config.size] size
* @param {object} [config.embeddings] embedding vectors
*/
constructor({ size = 512, embeddings = {}, ...rest }) {
super(rest)
this._size = size
this._v = embeddings
}
calc(x) {
this._i = x
const size = [...x.sizes, this._size]
const o = size.length === 2 ? new Matrix(...size) : new Tensor(size)
for (let i = 0; i < x.length; i++) {
if (!this._v[x.value[i]]) {
this._v[x.value[i]] = Matrix.randn(1, this._size)
}
for (let k = 0; k < this._size; k++) {
o.value[i * this._size + k] = this._v[x.value[i]].at(0, k)
}
}
return o
}
grad(bo) {
this._dw = {}
for (let i = 0; i < this._i.length; i++) {
for (let k = 0; k < this._size; k++) {
if (!this._dw[this._i.value[i]]) {
this._dw[this._i.value[i]] = Matrix.zeros(1, this._size)
}
this._dw[this._i.value[i]].add(
new Matrix(1, this._size, bo.value.slice(i * this._size, (i + 1) * this._size))
)
}
}
const bi = this._i.copy()
bi.fill(0)
return bi
}
update(optimizer) {
for (const w of Object.keys(this._dw)) {
this._v[w].sub(optimizer.delta(w, this._dw[w]))
}
}
toObject() {
return {
type: 'embedding',
size: this._size,
embeddings: this._v,
}
}
}
EmbeddingLayer.registLayer()