/
global_averagepool.js
99 lines (93 loc) · 2.58 KB
/
global_averagepool.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
96
97
98
99
import Layer, { NeuralnetworkLayerException } from './base.js'
import Tensor from '../../../util/tensor.js'
/**
* Global average pool layer
*/
export default class GlobalAveragePoolLayer extends Layer {
/**
* @param {object} config object
* @param {number} [config.channel_dim] Dimension of the channel
*/
constructor({ channel_dim = -1, ...rest }) {
super(rest)
this._channel_dim = channel_dim
if (this._channel_dim !== -1 && this._channel_dim !== 1) {
throw new NeuralnetworkLayerException('Invalid channel dimension.')
}
}
_index(i, c, k) {
return this._channel_dim === -1 ? [i, ...k, c] : [i, c, ...k]
}
calc(x) {
this._i = x
const traceSize = x.sizes.slice(1)
const outSize = Array(x.sizes.length).fill(1)
outSize[0] = x.sizes[0]
if (this._channel_dim === -1) {
outSize[x.dimension - 1] = x.sizes[x.dimension - 1]
traceSize.pop()
} else if (this._channel_dim === 1) {
outSize[1] = x.sizes[1]
traceSize.splice(0, 1)
}
const channels = this._channel_dim === -1 ? x.sizes[x.dimension - 1] : x.sizes[1]
this._o = new Tensor(outSize)
const count = this._i.length / this._o.length
for (let i = 0; i < x.sizes[0]; i++) {
for (let c = 0; c < channels; c++) {
const idx = Array(x.dimension - 2).fill(0)
let sumval = 0
do {
sumval += x.at(this._index(i, c, idx))
for (let k = 0; k < idx.length; k++) {
idx[k]++
if (idx[k] < traceSize[k]) {
break
}
idx[k] = 0
}
} while (idx.some(v => v > 0))
this._o.set(this._index(i, c, Array(x.dimension - 2).fill(0)), sumval / count)
}
}
return this._o
}
grad(bo) {
this._bo = bo
this._bi = new Tensor(this._i.sizes)
const traceSize = this._i.sizes.slice(1)
if (this._channel_dim === -1) {
traceSize.pop()
} else if (this._channel_dim === 1) {
traceSize.splice(0, 1)
}
const channels = this._channel_dim === -1 ? this._i.sizes[this._i.dimension - 1] : this._i.sizes[1]
const count = this._i.length / this._o.length
for (let i = 0; i < this._i.sizes[0]; i++) {
for (let c = 0; c < channels; c++) {
const idx = Array(this._i.dimension - 2).fill(0)
do {
this._bi.set(
this._index(i, c, idx),
this._bo.at(this._index(i, c, Array(this._i.dimension - 2).fill(0))) / count
)
for (let k = 0; k < idx.length; k++) {
idx[k]++
if (idx[k] < traceSize[k]) {
break
}
idx[k] = 0
}
} while (idx.some(v => v > 0))
}
}
return this._bi
}
toObject() {
return {
type: 'global_average_pool',
channel_dim: this._channel_dim,
}
}
}
GlobalAveragePoolLayer.registLayer()