/
prelu.js
95 lines (88 loc) · 2.12 KB
/
prelu.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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
import Layer from './base.js'
import Tensor from '../../../util/tensor.js'
/**
* Parametric ReLU layer
*/
export default class ParametricReLULayer extends Layer {
/**
* @param {object} config config
* @param {number | number[] | string} [config.a] a
*/
constructor({ a = 0.25, ...rest }) {
super(rest)
this._a = null
this._da = 0
if (typeof a === 'string') {
this._aname = a
} else if (Array.isArray(a)) {
this._a = Tensor.fromArray(a)
this._da = this._a.copy()
this._da.fill(0)
} else {
this._a = a
}
}
calc(x) {
this._i = x
this._o = x.copy()
if (this._aname) {
this._a = this.graph.getNode(this._aname).outputValue
}
this._o.broadcastOperate(this._a, (a, b) => (a > 0 ? a : b * a))
return this._o
}
grad(bo) {
const bi = this._i.copy()
bi.broadcastOperate(this._a, (a, b) => (a > 0 ? 1 : b))
bi.broadcastOperate(bo, (a, b) => a * b)
if (typeof this._a === 'number') {
this._da0 = 0
for (let i = 0; i < this._i.length; i++) {
if (this._i.value[i] < 0) {
this._da0 += bo.value[i] * this._i.value[i]
}
}
} else {
this._da0 = this._a.copy()
this._da0.fill(0)
const dimdiff = this._i.dimension - this._a.dimension
const idx = Array(this._i.dimension).fill(0)
do {
const val = this._i.at(idx)
if (val > 0) {
this._da0.operateAt(idx.slice(dimdiff), v => v + val * bo.at(idx))
}
for (let k = 0; k < idx.length; k++) {
idx[k]++
if (idx[k] < this._i.sizes[k]) {
break
}
idx[k] = 0
}
} while (idx.some(v => v > 0))
}
if (this._aname) {
return [bi, { [this._aname]: this._da0 }]
}
return bi
}
update(optimizer) {
const myu = 0.1
if (this._aname) {
return
} else if (typeof this._a === 'number') {
this._da = myu * this._da + (optimizer.lr * this._da0) / this._i.length
this._a -= this._da
} else {
this._da.broadcastOperate(this._da0, (a, b) => myu * a + (optimizer.lr * b) / this._i.length)
this._a.broadcastOperate(this._da, (a, b) => a - b)
}
}
toObject() {
return {
type: 'prelu',
a: this._aname || this._a,
}
}
}
ParametricReLULayer.registLayer('prelu')