Add UpSampling layer (#972)

Author: ishii-norimi

README.md

@@ -213,7 +213,7 @@ for (let i = 0; i < n; i++) {
 | function | absolute, acos, acosh, APL, Aranda, asin, asinh, atan, atanh, attention, batch normalization, BDAA, Bent identity, BLU, BReLU, ceil, CELU, cloglog, cloglogm, cos, cosh, CReLU, EELU, (hard) ELiSH, Elliott, ELU, embedding, EReLU, erf, ESwish, exp, FELU, full, floor, FReLU, gaussian, GELU, Hard shrink, Hexpo, identity, ISigmoid, layer normalization, Leaky ReLU, LiSHT, log, loglog, logsigmoid, mish, MPELU, MTLU, negative, NLReLU, PAU, PDELU, PELU, PLU, PReLU, PREU, PSF, pTanh, PTELU, reciprocal, ReLU, RePU, ReSech, REU, rootsig, round, RReLU, RTReLU, SELU, (hard) sigmoid, sign, SiLU, sin, sinh, SLAF, SLU, softmax, softplus, Soft shrink, softsign, sqrt, square, SReLU, SRS, sSigmoid, sTanh, (hard) Swish, TAF, tan, (hard) tanh, tanhExp, tanShrink, Thresholded ReLU |
 | operator | add, sub, mult, div, mod, matmul, power, max, min |
 | logical | and, bitwise and, bitwise not, bitwise or, bitwise xor, equal, greater, greater or equal, is inf, is nan, left bitshift, less, less or equal, not, or, right bitshift, xor |
-| convolute | convolution, (Global) MaxPool, (Global) AveragePool, (Global) LpPool, LRN |
+| convolute | convolution, (Global) MaxPool, (Global) AveragePool, (Global) LpPool, LRN, UpSampling |
 | recurrent | GRU, LSTM, Simple RNN |
 | reduce | sum, mean, prod, variance, std, reduce max, reduce min, argmax, softargmax |
 | graph | convolutional, SAGE, readout |

lib/model/nns/layer/conv.js

@@ -10,7 +10,7 @@ export default class ConvLayer extends Layer {
 	 * @param {number | number[]} config.kernel Size of kernel
 	 * @param {number} [config.channel] Number of output channel
 	 * @param {number | number[]} [config.stride] Step of stride
-	 * @param {number | number[]} [config.padding] size of padding
+	 * @param {number | number[] | [number, number][]} [config.padding] size of padding
 	 * @param {number[][] | Tensor | string} [config.w] Weight of kernel
 	 * @param {string | object} [config.activation] Name of activation or activation layer object
 	 * @param {number} [config.l2_decay] L2 decay

lib/model/nns/layer/index.js

@@ -99,6 +99,7 @@ export { default as SwishLayer } from './swish.js'
 export { default as TrainableAFLayer } from './taf.js'
 export { default as ThresholdedReLULayer } from './thresholded_relu.js'
 export { default as TransposeLayer } from './transpose.js'
+export { default as UpSamplingLayer } from './upsampling.js'
 export { default as VariableLayer } from './variable.js'
 export { default as VarLayer } from './variance.js'
 
@@ -145,7 +146,7 @@ export { default as VarLayer } from './variance.js'
  * { type: 'concat', axis?: number } |
  * { type: 'cond' } |
  * { type: 'const', value: number } |
- * { type: 'conv', kernel: number | number[], channel?: number, stride?: number | number[], padding?: number | number[], w?: number[][] | Tensor | string, activation?: string | object, l2_decay?: number, l1_decay?: number, channel_dim?: number } |
+ * { type: 'conv', kernel: number | number[], channel?: number, stride?: number | number[], padding?: number | number[] | [number, number][], w?: number[][] | Tensor | string, activation?: string | object, l2_decay?: number, l1_decay?: number, channel_dim?: number } |
  * { type: 'cos' } |
  * { type: 'cosh' } |
  * { type: 'crelu' } |
@@ -283,6 +284,7 @@ export { default as VarLayer } from './variance.js'
  * { type: 'tanhshrink' } |
  * { type: 'thresholded_relu', a?: number } |
  * { type: 'transpose', axis: number[] } |
+ * { type: 'up_sampling', size: number | number[], channel_dim?: number } |
  * { type: 'variable', size: number[] | string, l2_decay?: number, l1_decay?: number, value?: number[] | number[][] | Tensor } |
  * { type: 'variance', axis?: number | number[] | string, keepdims?: boolean } |
  * { type: 'xor' }

lib/model/nns/layer/upsampling.js

@@ -0,0 +1,118 @@
+import Layer, { NeuralnetworkLayerException } from './base.js'
+import Tensor from '../../../util/tensor.js'
+
+/**
+ * Max pool layer
+ */
+export default class UpSamplingLayer extends Layer {
+	/**
+	 * @param {object} config object
+	 * @param {number | number[]} config.size Size of kernel
+	 * @param {number} [config.channel_dim] Dimension of the channel
+	 */
+	constructor({ size, channel_dim = -1, ...rest }) {
+		super(rest)
+		this._size = size
+		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) {
+		if (!Array.isArray(this._size)) {
+			this._size = Array(x.dimension - 2).fill(this._size)
+		}
+		if (x.dimension !== this._size.length + 2) {
+			throw new NeuralnetworkLayerException('Invalid size', [this, x])
+		}
+		this._i = x
+		const koff = this._channel_dim === -1 ? 1 : 2
+		const outSize = [x.sizes[0], ...this._size.map((k, d) => x.sizes[d + koff] * k)]
+		if (this._channel_dim === -1) {
+			outSize.push(x.sizes[x.dimension - 1])
+		} else if (this._channel_dim === 1) {
+			outSize.splice(1, 0, x.sizes[1])
+		}
+		const channels = this._channel_dim === -1 ? x.sizes[x.dimension - 1] : x.sizes[1]
+		this._o = new Tensor(outSize)
+		for (let i = 0; i < x.sizes[0]; i++) {
+			for (let c = 0; c < channels; c++) {
+				const idx = Array(x.dimension - 2).fill(0)
+				do {
+					const offset = Array(x.dimension - 2).fill(0)
+					do {
+						const p = idx.map((v, i) => v * this._size[i] + offset[i])
+						this._o.set(this._index(i, c, p), x.at(this._index(i, c, idx)))
+						for (let k = 0; k < offset.length; k++) {
+							offset[k]++
+							if (offset[k] < this._size[k]) {
+								break
+							}
+							offset[k] = 0
+						}
+					} while (offset.some(v => v > 0))
+					for (let k = 0; k < idx.length; k++) {
+						idx[k]++
+						if (idx[k] < this._i.sizes[k + koff]) {
+							break
+						}
+						idx[k] = 0
+					}
+				} while (idx.some(v => v > 0))
+			}
+		}
+		return this._o
+	}
+
+	grad(bo) {
+		this._bo = bo
+		this._bi = new Tensor(this._i.sizes)
+		const koff = this._channel_dim === -1 ? 1 : 2
+		const channels = this._channel_dim === -1 ? this._i.sizes[this._i.dimension - 1] : this._i.sizes[1]
+		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 {
+					const offset = Array(this._i.dimension - 2).fill(0)
+					let sum = 0
+					do {
+						const p = idx.map((v, i) => v * this._size[i] + offset[i])
+						sum += this._bo.at(this._index(i, c, p))
+						for (let k = 0; k < offset.length; k++) {
+							offset[k]++
+							if (offset[k] < this._size[k]) {
+								break
+							}
+							offset[k] = 0
+						}
+					} while (offset.some(v => v > 0))
+
+					this._bi.set(this._index(i, c, idx), sum)
+					for (let k = 0; k < idx.length; k++) {
+						idx[k]++
+						if (idx[k] < this._i.sizes[k + koff]) {
+							break
+						}
+						idx[k] = 0
+					}
+				} while (idx.some(v => v > 0))
+			}
+		}
+		return this._bi
+	}
+
+	toObject() {
+		return {
+			type: 'up_sampling',
+			size: this._size,
+			channel_dim: this._channel_dim,
+		}
+	}
+}
+
+UpSamplingLayer.registLayer()

lib/model/nns/onnx/layer/index.js

@@ -146,5 +146,6 @@ export { default as tanhexp } from './tanhexp.js'
 export { default as tanhshrink } from './tanhshrink.js'
 export { default as thresholded_relu } from './thresholded_relu.js'
 export { default as transpose } from './transpose.js'
+export { default as up_sampling } from './up_sampling.js'
 export { default as variance } from './variance.js'
 export { default as xor } from './xor.js'

lib/model/nns/onnx/layer/up_sampling.js

@@ -0,0 +1,49 @@
+import { onnx } from '../onnx_exporter.js'
+
+/**
+ * Handle upsampling layer
+ */
+export default {
+	/**
+	 * Export to onnx object.
+	 * @param {onnx.ModelProto} model Model object
+	 * @param {import("../../graph").LayerObject & {type: 'up_sampling'}} obj Node object
+	 * @param {{[key: string]: {type: onnx.TensorProto.DataType; size: number[]}}} info Output informatino of other layers
+	 * @returns {{type: onnx.TensorProto.DataType; size: number[]}} Output information of this layer
+	 */
+	export(model, obj, info) {
+		const input = Array.isArray(obj.input) ? obj.input[0] : obj.input
+		const inSize = info[input].size
+		const scale = Array.isArray(obj.size) ? obj.size : Array(inSize.length - 2).fill(obj.size)
+		scale.unshift(1)
+		if (obj.channel_dim == null || obj.channel_dim === -1) {
+			scale.push(1)
+		} else if (obj.channel_dim === 1) {
+			scale.splice(1, 0, 1)
+		}
+		const outSize = inSize.map((v, i) => (v == null ? null : v * scale[i]))
+
+		const tensor_scale = new onnx.TensorProto()
+		tensor_scale.setName(obj.name + '_scale')
+		tensor_scale.setDataType(onnx.TensorProto.DataType.FLOAT)
+		tensor_scale.setDimsList([scale.length])
+		tensor_scale.setFloatDataList(scale)
+
+		const node = new onnx.NodeProto()
+		node.setOpType('Resize')
+		node.addInput(input)
+		node.addInput('')
+		node.addInput(obj.name + '_scale')
+		node.addOutput(obj.name)
+		const mode = new onnx.AttributeProto()
+		mode.setName('mode')
+		mode.setType(onnx.AttributeProto.AttributeType.STRING)
+		mode.setS(new TextEncoder().encode('nearest'))
+		node.addAttribute(mode)
+
+		const graph = model.getGraph()
+		graph.addInitializer(tensor_scale)
+		graph.addNode(node)
+		return { size: outSize }
+	},
+}