Add ONNX export (#928)

* Add ONNX export

* Fix some issues

* Fix test

* set env.wasm.numThreads to 1

Author: ishii-norimi

create_import_list.js

@@ -223,6 +223,19 @@ const createONNXOperatorlist = async () => {
 	await fs.promises.writeFile(operatorsDir + '/index.js', '// This file is generated automatically.\n' + code)
 }
 
+const createONNXLayerlist = async () => {
+	const layerDir = './lib/model/nns/onnx/layer'
+	const files = await fs.promises.readdir(layerDir)
+	let code = ''
+	for (const file of files) {
+		if (file !== 'index.js' && file.endsWith('.js')) {
+			code += `export { default as ${file.slice(0, -3)} } from './${file}'\n`
+		}
+	}
+	await fs.promises.writeFile(layerDir + '/index.js', '// This file is generated automatically.\n' + code)
+}
+
 await createLayerlist()
+await createONNXLayerlist()
 await createONNXOperatorlist()
 await createEntrypoint()

lib/model/nns/graph.js

@@ -2,6 +2,7 @@ import Matrix from '../../util/matrix.js'
 import { NeuralnetworkException } from '../neuralnetwork.js'
 import Layer from './layer/base.js'
 import { InputLayer, OutputLayer } from './layer/index.js'
+import ONNXExporter from './onnx/onnx_exporter.js'
 import ONNXImporter from './onnx/onnx_importer.js'
 
 /**
@@ -130,6 +131,14 @@ export default class ComputationalGraph {
 		return s + '}'
 	}
 
+	/**
+	 * Returns onnx model
+	 * @returns {Uint8Array} onnx model byte array
+	 */
+	toONNX() {
+		return ONNXExporter.dump(this.toObject())
+	}
+
 	/**
 	 * Add a layer.
 	 * @param {Layer | PlainLayerObject} layer Added layer

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

@@ -0,0 +1,22 @@
+import { onnx } from '../onnx_exporter.js'
+
+/**
+ * Handle abs layer
+ */
+export default {
+	/**
+	 * Export to onnx object.
+	 * @param {onnx.ModelProto} model Model object
+	 * @param {import("../../graph").LayerObject & {type: 'abs'}} obj Node object
+	 */
+	export(model, obj) {
+		const input = Array.isArray(obj.input) ? obj.input[0] : obj.input
+		const node = new onnx.NodeProto()
+		node.setOpType('Abs')
+		node.addInput(input)
+		node.addOutput(obj.name)
+
+		const graph = model.getGraph()
+		graph.addNode(node)
+	},
+}

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

@@ -0,0 +1,22 @@
+import { onnx } from '../onnx_exporter.js'
+
+/**
+ * Handle acos layer
+ */
+export default {
+	/**
+	 * Export to onnx object.
+	 * @param {onnx.ModelProto} model Model object
+	 * @param {import("../../graph").LayerObject & {type: 'acos'}} obj Node object
+	 */
+	export(model, obj) {
+		const input = Array.isArray(obj.input) ? obj.input[0] : obj.input
+		const node = new onnx.NodeProto()
+		node.setOpType('Acos')
+		node.addInput(input)
+		node.addOutput(obj.name)
+
+		const graph = model.getGraph()
+		graph.addNode(node)
+	},
+}

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

@@ -0,0 +1,22 @@
+import { onnx } from '../onnx_exporter.js'
+
+/**
+ * Handle acosh layer
+ */
+export default {
+	/**
+	 * Export to onnx object.
+	 * @param {onnx.ModelProto} model Model object
+	 * @param {import("../../graph").LayerObject & {type: 'acosh'}} obj Node object
+	 */
+	export(model, obj) {
+		const input = Array.isArray(obj.input) ? obj.input[0] : obj.input
+		const node = new onnx.NodeProto()
+		node.setOpType('Acosh')
+		node.addInput(input)
+		node.addOutput(obj.name)
+
+		const graph = model.getGraph()
+		graph.addNode(node)
+	},
+}

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

@@ -0,0 +1,26 @@
+import { onnx } from '../onnx_exporter.js'
+
+/**
+ * Handle add layer
+ */
+export default {
+	/**
+	 * Export to onnx object.
+	 * @param {onnx.ModelProto} model Model object
+	 * @param {import("../../graph").LayerObject & {type: 'add'}} obj Node object
+	 */
+	export(model, obj) {
+		if (!Array.isArray(obj.input)) {
+			throw new Error(`Invalid attribute 'input' value ${obj.input}.`)
+		}
+		const node = new onnx.NodeProto()
+		node.setOpType(obj.input.length === 2 ? 'Add' : 'Sum')
+		for (const i of obj.input) {
+			node.addInput(i)
+		}
+		node.addOutput(obj.name)
+
+		const graph = model.getGraph()
+		graph.addNode(node)
+	},
+}

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

@@ -0,0 +1,64 @@
+import { onnx } from '../onnx_exporter.js'
+
+/**
+ * Handle and layer
+ */
+export default {
+	/**
+	 * Export to onnx object.
+	 * @param {onnx.ModelProto} model Model object
+	 * @param {import("../../graph").LayerObject & {type: 'and'}} 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[]} | undefined} Output information of this layer
+	 */
+	export(model, obj, info) {
+		if (!Array.isArray(obj.input)) {
+			throw new Error(`Invalid attribute 'input' value ${obj.input}.`)
+		}
+		const graph = model.getGraph()
+
+		const node = new onnx.NodeProto()
+		if (obj.input.length === 1) {
+			node.setOpType('Identity')
+			node.addInput(obj.input[0])
+			node.addOutput(obj.name)
+			graph.addNode(node)
+			return
+		}
+		const boolInputs = []
+		for (const i of obj.input) {
+			if (info[i].type === onnx.TensorProto.DataType.BOOL) {
+				boolInputs.push(i)
+			} else {
+				const castnode = new onnx.NodeProto()
+				castnode.setOpType('Cast')
+				castnode.addInput(i)
+				castnode.addOutput(`${obj.name}_${i}_cast`)
+				const to = new onnx.AttributeProto()
+				to.setName('to')
+				to.setType(onnx.AttributeProto.AttributeType.INT)
+				to.setI(onnx.TensorProto.DataType.BOOL)
+				castnode.addAttribute(to)
+				graph.addNode(castnode)
+				boolInputs.push(`${obj.name}_${i}_cast`)
+			}
+		}
+		let prev_in = boolInputs[0]
+		for (let i = 1; i < boolInputs.length - 1; i++) {
+			const node_and = new onnx.NodeProto()
+			node_and.setOpType('And')
+			node_and.addInput(prev_in)
+			node_and.addInput(boolInputs[i])
+			node_and.addOutput((prev_in = obj.name + `_and_${i - 1}`))
+			graph.addNode(node_and)
+		}
+
+		node.setOpType('And')
+		node.addInput(prev_in)
+		node.addInput(boolInputs.at(-1))
+
+		node.addOutput(obj.name)
+		graph.addNode(node)
+		return { type: onnx.TensorProto.DataType.BOOL }
+	},
+}

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

@@ -0,0 +1,71 @@
+import { onnx } from '../onnx_exporter.js'
+
+/**
+ * Handle apl layer
+ */
+export default {
+	/**
+	 * Export to onnx object.
+	 * @param {onnx.ModelProto} model Model object
+	 * @param {import("../../graph").LayerObject & {type: 'apl'}} obj Node object
+	 */
+	export(model, obj) {
+		const input = Array.isArray(obj.input) ? obj.input[0] : obj.input
+		const node_base = new onnx.NodeProto()
+		node_base.setOpType('Relu')
+		node_base.addInput(input)
+		node_base.addOutput(obj.name + '_base')
+
+		const graph = model.getGraph()
+		graph.addNode(node_base)
+
+		const node_add = new onnx.NodeProto()
+		node_add.setOpType('Sum')
+		node_add.addInput(obj.name + '_base')
+
+		const s = obj.s ?? 2
+		const a = Array.isArray(obj.a) ? obj.a : Array(s).fill(obj.a ?? 0.1)
+		const b = Array.isArray(obj.b) ? obj.b : Array(s).fill(obj.b ?? 0)
+		for (let i = 0; i < s; i++) {
+			const tensor_a = new onnx.TensorProto()
+			tensor_a.setName(obj.name + '_a_' + i)
+			tensor_a.setDataType(onnx.TensorProto.DataType.FLOAT)
+			tensor_a.setDimsList([1])
+			tensor_a.setFloatDataList([a[i]])
+
+			const tensor_b = new onnx.TensorProto()
+			tensor_b.setName(obj.name + '_b_' + i)
+			tensor_b.setDataType(onnx.TensorProto.DataType.FLOAT)
+			tensor_b.setDimsList([1])
+			tensor_b.setFloatDataList([b[i]])
+
+			const node_sub = new onnx.NodeProto()
+			node_sub.setOpType('Sub')
+			node_sub.addInput(obj.name + '_b_' + i)
+			node_sub.addInput(input)
+			node_sub.addOutput(obj.name + '_sub_' + i)
+
+			const node_relu = new onnx.NodeProto()
+			node_relu.setOpType('Relu')
+			node_relu.addInput(obj.name + '_sub_' + i)
+			node_relu.addOutput(obj.name + '_relu_' + i)
+
+			const node_mult = new onnx.NodeProto()
+			node_mult.setOpType('Mul')
+			node_mult.addInput(obj.name + '_relu_' + i)
+			node_mult.addInput(obj.name + '_a_' + i)
+			node_mult.addOutput(obj.name + '_mul_' + i)
+
+			node_add.addInput(obj.name + '_mul_' + i)
+
+			graph.addInitializer(tensor_a)
+			graph.addInitializer(tensor_b)
+			graph.addNode(node_sub)
+			graph.addNode(node_relu)
+			graph.addNode(node_mult)
+		}
+
+		node_add.addOutput(obj.name)
+		graph.addNode(node_add)
+	},
+}

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

@@ -0,0 +1,72 @@
+import { onnx } from '../onnx_exporter.js'
+import { getConstNodeName } from '../utils.js'
+
+/**
+ * Handle aranda layer
+ */
+export default {
+	/**
+	 * Export to onnx object.
+	 * @param {onnx.ModelProto} model Model object
+	 * @param {import("../../graph").LayerObject & {type: 'aranda'}} obj Node object
+	 */
+	export(model, obj) {
+		const input = Array.isArray(obj.input) ? obj.input[0] : obj.input
+		const node_exp = new onnx.NodeProto()
+		node_exp.setOpType('Exp')
+		node_exp.addInput(input)
+		node_exp.addOutput(obj.name + '_exp')
+
+		const tensor_l = new onnx.TensorProto()
+		tensor_l.setName(obj.name + '_l')
+		tensor_l.setDataType(onnx.TensorProto.DataType.FLOAT)
+		tensor_l.setDimsList([1])
+		tensor_l.setFloatDataList([obj.l ?? 2])
+
+		const tensor1 = getConstNodeName(model, 1)
+
+		const node_mult = new onnx.NodeProto()
+		node_mult.setOpType('Mul')
+		node_mult.addInput(obj.name + '_exp')
+		node_mult.addInput(obj.name + '_l')
+		node_mult.addOutput(obj.name + '_mult')
+
+		const node_add = new onnx.NodeProto()
+		node_add.setOpType('Add')
+		node_add.addInput(obj.name + '_mult')
+		node_add.addInput(tensor1)
+		node_add.addOutput(obj.name + '_add')
+
+		const node_reciprocal = new onnx.NodeProto()
+		node_reciprocal.setOpType('Reciprocal')
+		node_reciprocal.addInput(obj.name + '_add')
+		node_reciprocal.addOutput(obj.name + '_reciprocal')
+
+		const node_reciprocal_pow = new onnx.NodeProto()
+		node_reciprocal_pow.setOpType('Reciprocal')
+		node_reciprocal_pow.addInput(obj.name + '_l')
+		node_reciprocal_pow.addOutput(obj.name + '_reciprocal_pow')
+
+		const node_pow = new onnx.NodeProto()
+		node_pow.setOpType('Pow')
+		node_pow.addInput(obj.name + '_reciprocal')
+		node_pow.addInput(obj.name + '_reciprocal_pow')
+		node_pow.addOutput(obj.name + '_pow')
+
+		const node_sub = new onnx.NodeProto()
+		node_sub.setOpType('Sub')
+		node_sub.addInput(tensor1)
+		node_sub.addInput(obj.name + '_pow')
+		node_sub.addOutput(obj.name)
+
+		const graph = model.getGraph()
+		graph.addInitializer(tensor_l)
+		graph.addNode(node_exp)
+		graph.addNode(node_mult)
+		graph.addNode(node_add)
+		graph.addNode(node_reciprocal)
+		graph.addNode(node_reciprocal_pow)
+		graph.addNode(node_pow)
+		graph.addNode(node_sub)
+	},
+}

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

@@ -0,0 +1,45 @@
+import { onnx } from '../onnx_exporter.js'
+
+/**
+ * Handle argmax layer
+ */
+export default {
+	/**
+	 * Export to onnx object.
+	 * @param {onnx.ModelProto} model Model object
+	 * @param {import("../../graph").LayerObject & {type: 'argmax'}} 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 graph = model.getGraph()
+
+		const input = Array.isArray(obj.input) ? obj.input[0] : obj.input
+		const node = new onnx.NodeProto()
+		node.setOpType('ArgMax')
+		node.addInput(input)
+		node.addOutput(obj.name)
+
+		const axis = new onnx.AttributeProto()
+		axis.setName('axis')
+		axis.setType(onnx.AttributeProto.AttributeType.INT)
+		axis.setI(obj.axis ?? -1)
+		node.addAttribute(axis)
+		const keepdims = new onnx.AttributeProto()
+		keepdims.setName('keepdims')
+		keepdims.setType(onnx.AttributeProto.AttributeType.INT)
+		keepdims.setI(obj.keepdims ?? true ? 1 : 0)
+		node.addAttribute(keepdims)
+
+		graph.addNode(node)
+
+		const size = info[input].size.concat()
+		const targetAxis = axis.getI() < 0 ? axis.getI() + size.length : axis.getI()
+		if (obj.keepdims ?? true) {
+			size[targetAxis] = 1
+		} else {
+			size.splice(targetAxis, 1)
+		}
+		return { type: onnx.TensorProto.DataType.INT64, size }
+	},
+}