Add as constant layer if input is initializer (#934)

Author: ishii-norimi

lib/model/nns/onnx/onnx_importer.js

@@ -5,6 +5,7 @@ import input from './operators/input.js'
 import output from './operators/output.js'
 
 import * as operators from './operators/index.js'
+import { loadTensor } from './utils.js'
 
 /**
  * ONNX importer
@@ -41,6 +42,22 @@ export default class ONNXImporter {
 		for (const node of graph.getOutputList()) {
 			nodes.push(...output.import(model, node))
 		}
+
+		const importedNames = new Set(nodes.map(n => n.name))
+		const inputNames = new Map(graph.getInitializerList().map(init => [init.getName(), init]))
+		for (const node of nodes.filter(n => n.input)) {
+			const inputs = Array.isArray(node.input) ? node.input : [node.input]
+			for (const i of inputs) {
+				if (importedNames.has(i)) {
+					continue
+				}
+				const initializer = inputNames.get(i)
+				if (initializer) {
+					nodes.push({ type: 'const', name: i, value: loadTensor(initializer) })
+					importedNames.add(i)
+				}
+			}
+		}
 		return nodes
 	}
 }

lib/model/nns/onnx/operators/add.js

@@ -1,5 +1,4 @@
 import { onnx } from '../onnx_importer.js'
-import { requireTensor } from '../utils.js'
 
 /**
  * Handle add operator
@@ -14,9 +13,6 @@ export default {
 	 * @returns {object[]} Objects represented a layer
 	 */
 	import(model, node) {
-		return [
-			...requireTensor(model, node.getInputList()),
-			{ type: 'add', input: node.getInputList(), name: node.getOutputList()[0] },
-		]
+		return [{ type: 'add', input: node.getInputList(), name: node.getOutputList()[0] }]
 	},
 }

lib/model/nns/onnx/operators/div.js

@@ -1,5 +1,4 @@
 import { onnx } from '../onnx_importer.js'
-import { requireTensor } from '../utils.js'
 
 /**
  * Handle div operator
@@ -14,9 +13,6 @@ export default {
 	 * @returns {object[]} Objects represented a layer
 	 */
 	import(model, node) {
-		return [
-			...requireTensor(model, node.getInputList()),
-			{ type: 'div', input: node.getInputList(), name: node.getOutputList()[0] },
-		]
+		return [{ type: 'div', input: node.getInputList(), name: node.getOutputList()[0] }]
 	},
 }

lib/model/nns/onnx/operators/matmul.js

@@ -1,5 +1,4 @@
 import { onnx } from '../onnx_importer.js'
-import { requireTensor } from '../utils.js'
 
 /**
  * Handle matmul operator
@@ -14,9 +13,6 @@ export default {
 	 * @returns {object[]} Objects represented a layer
 	 */
 	import(model, node) {
-		return [
-			...requireTensor(model, node.getInputList()),
-			{ type: 'matmul', input: node.getInputList(), name: node.getOutputList()[0] },
-		]
+		return [{ type: 'matmul', input: node.getInputList(), name: node.getOutputList()[0] }]
 	},
 }

lib/model/nns/onnx/operators/max.js

@@ -1,5 +1,4 @@
 import { onnx } from '../onnx_importer.js'
-import { requireTensor } from '../utils.js'
 
 /**
  * Handle max operator
@@ -14,9 +13,6 @@ export default {
 	 * @returns {object[]} Objects represented a layer
 	 */
 	import(model, node) {
-		return [
-			...requireTensor(model, node.getInputList()),
-			{ type: 'max', input: node.getInputList(), name: node.getOutputList()[0] },
-		]
+		return [{ type: 'max', input: node.getInputList(), name: node.getOutputList()[0] }]
 	},
 }

lib/model/nns/onnx/operators/min.js

@@ -1,5 +1,4 @@
 import { onnx } from '../onnx_importer.js'
-import { requireTensor } from '../utils.js'
 
 /**
  * Handle min operator
@@ -14,9 +13,6 @@ export default {
 	 * @returns {object[]} Objects represented a layer
 	 */
 	import(model, node) {
-		return [
-			...requireTensor(model, node.getInputList()),
-			{ type: 'min', input: node.getInputList(), name: node.getOutputList()[0] },
-		]
+		return [{ type: 'min', input: node.getInputList(), name: node.getOutputList()[0] }]
 	},
 }

lib/model/nns/onnx/operators/mul.js

@@ -1,5 +1,4 @@
 import { onnx } from '../onnx_importer.js'
-import { requireTensor } from '../utils.js'
 
 /**
  * Handle mul operator
@@ -14,9 +13,6 @@ export default {
 	 * @returns {object[]} Objects represented a layer
 	 */
 	import(model, node) {
-		return [
-			...requireTensor(model, node.getInputList()),
-			{ type: 'mult', input: node.getInputList(), name: node.getOutputList()[0] },
-		]
+		return [{ type: 'mult', input: node.getInputList(), name: node.getOutputList()[0] }]
 	},
 }

lib/model/nns/onnx/operators/reshape.js

@@ -1,5 +1,5 @@
 import { onnx } from '../onnx_importer.js'
-import { loadTensor, loadAttribute, requireTensor } from '../utils.js'
+import { loadTensor, loadAttribute } from '../utils.js'
 
 /**
  * Handle reshape operator
@@ -31,7 +31,6 @@ export default {
 			throw new Error(`Invalid shape value ${JSON.stringify(initializers.shape)}.`)
 		}
 		return [
-			...requireTensor(model, inputList[0]),
 			{
 				type: 'reshape',
 				input: [inputList[0]],

lib/model/nns/onnx/operators/sub.js

@@ -1,5 +1,4 @@
 import { onnx } from '../onnx_importer.js'
-import { requireTensor } from '../utils.js'
 
 /**
  * Handle sub operator
@@ -14,9 +13,6 @@ export default {
 	 * @returns {object[]} Objects represented a layer
 	 */
 	import(model, node) {
-		return [
-			...requireTensor(model, node.getInputList()),
-			{ type: 'sub', input: node.getInputList(), name: node.getOutputList()[0] },
-		]
+		return [{ type: 'sub', input: node.getInputList(), name: node.getOutputList()[0] }]
 	},
 }

lib/model/nns/onnx/utils.js

@@ -118,22 +118,6 @@ export const loadAttribute = attribute => {
 	throw new Error('Not implemented attribute type.')
 }
 
-/**
- * Create const layers from initializer list.
- * @param {onnx.ModelProto} model Model object
- * @param {string[]} names Input name
- * @returns {object[]} Require layer objects
- */
-export const requireTensor = (model, names) => {
-	const layers = []
-	for (const initializer of model.getGraph().getInitializerList()) {
-		if (names.includes(initializer.getName())) {
-			layers.push({ type: 'const', value: loadTensor(initializer), name: initializer.getName() })
-		}
-	}
-	return layers
-}
-
 /**
  * Create const node if needed and return const node name.
  * @param {onnx.ModelProto} model Model object

tests/lib/model/nns/graph.test.js

@@ -11,6 +11,7 @@ import Matrix from '../../../../lib/util/matrix.js'
 import ComputationalGraph from '../../../../lib/model/nns/graph.js'
 
 import Layer from '../../../../lib/model/nns/layer/base.js'
+import Tensor from '../../../../lib/util/tensor.js'
 
 describe('Computational Graph', () => {
 	test('constructor', () => {
@@ -193,6 +194,40 @@ describe('Computational Graph', () => {
 				expect(y[i]).toBeCloseTo(Math.tanh(x.value[i]))
 			}
 		})
+
+		test('complex layers', async () => {
+			const graph = new ComputationalGraph()
+			graph.add(Layer.fromObject({ type: 'input', size: [null, 6, 6, 3] }))
+			graph.add(Layer.fromObject({ type: 'conv', kernel: 3 }))
+			graph.add(Layer.fromObject({ type: 'max_pool', kernel: 2 }))
+			graph.add(Layer.fromObject({ type: 'relu' }))
+			graph.add(Layer.fromObject({ type: 'flatten' }))
+			graph.add(Layer.fromObject({ type: 'full', out_size: 10 }))
+			graph.add(Layer.fromObject({ type: 'tanh' }), 'v')
+			graph.add(Layer.fromObject({ type: 'pau' }))
+			graph.add(Layer.fromObject({ type: 'tanh' }), 'pau')
+			graph.add(Layer.fromObject({ type: 'apl' }), 'apl', 'v')
+			graph.add(Layer.fromObject({ type: 'add' }), null, ['pau', 'apl'])
+			graph.add(Layer.fromObject({ type: 'output' }))
+
+			const x = Tensor.randn([100, 6, 6, 3])
+			graph.bind({ input: x })
+			graph.calc()
+			const t = graph.outputNodes[0].outputValue
+
+			const buf = await graph.toONNX()
+			session = await ort.InferenceSession.create(buf)
+
+			const xten = new ort.Tensor('float32', x.value, x.sizes)
+			const out = await session.run({ _input: xten })
+			const yten = out._add
+			expect(yten.dims).toEqual([100, 10])
+			const y = await yten.getData(true)
+
+			for (let i = 0; i < y.length; i++) {
+				expect(y[i]).toBeCloseTo(t.value[i])
+			}
+		})
 	})
 
 	describe('add', () => {