Skip outside calculation in maxpool layer (#691)

ai-on-browser · Nov 18, 2023 · 76a56b3 · 76a56b3
1 parent fc93ce1
commit 76a56b3
Showing 1 changed file with 13 additions and 22 deletions.
diff --git a/lib/model/nns/layer/maxpool.js b/lib/model/nns/layer/maxpool.js
@@ -73,15 +73,12 @@ export default class MaxPoolLayer extends Layer {
 					const offset = Array(x.dimension - 2).fill(0)
 					let maxval = -Infinity
 					do {
-						const p = idx.map((v, i) =>
-							Math.max(
-								0,
-								Math.min(x.sizes[i + koff] - 1, v * this._stride[i] - this._padding[i][0] + offset[i])
-							)
-						)
-						const v = x.at(this._index(i, c, p))
-						if (maxval < v) {
-							maxval = v
+						const p = idx.map((v, i) => v * this._stride[i] - this._padding[i][0] + offset[i])
+						if (p.every((v, i) => 0 <= v && v < x.sizes[i + koff])) {
+							const v = x.at(this._index(i, c, p))
+							if (maxval < v) {
+								maxval = v
+							}
 						}
 						for (let k = 0; k < offset.length; k++) {
 							offset[k]++
@@ -118,19 +115,13 @@ export default class MaxPoolLayer extends Layer {
 					let maxval = -Infinity
 					let maxidx = null
 					do {
-						const p = idx.map((v, i) =>
-							Math.max(
-								0,
-								Math.min(
-									this._i.sizes[i + koff] - 1,
-									v * this._stride[i] - this._padding[i][0] + offset[i]
-								)
-							)
-						)
-						const v = this._i.at(this._index(i, c, p))
-						if (maxval < v) {
-							maxval = v
-							maxidx = p
+						const p = idx.map((v, i) => v * this._stride[i] - this._padding[i][0] + offset[i])
+						if (p.every((v, i) => 0 <= v && v < this._i.sizes[i + koff])) {
+							const v = this._i.at(this._index(i, c, p))
+							if (maxval < v) {
+								maxval = v
+								maxidx = p
+							}
 						}
 						for (let k = 0; k < offset.length; k++) {
 							offset[k]++