Enhance ONNX model creation by adding configurable opset version parameter (#1037)

Author: ishii-norimi

lib/model/nns/onnx/onnx_exporter.js

@@ -10,17 +10,20 @@ import * as layers from './layer/index.js'
 export default class ONNXExporter {
 	/**
 	 * Create onnx model proto.
+	 * @param {object} [config] object
+	 * @param {object} [config.opset] Config for operator set
+	 * @param {number} [config.opset.version] Version of operator set
 	 * @returns {onnx.ModelProto} Model proto
 	 */
-	static createONNXModel() {
+	static createONNXModel({ opset } = {}) {
 		const model = new onnx.ModelProto()
 		model.setProducerName('ai-on-browser/data-analysis-models')
 		model.setProducerVersion('0.24.0')
 		model.setIrVersion(9)
 
 		const opsetImport = new onnx.OperatorSetIdProto()
 		opsetImport.setDomain('')
-		opsetImport.setVersion(19)
+		opsetImport.setVersion(opset?.version ?? 19)
 		model.addOpsetImport(opsetImport)
 
 		const graph = new onnx.GraphProto()

tests/lib/model/nns/onnx/layer/gelu.test.js

@@ -52,8 +52,7 @@ describe('export', () => {
 
 	describe('opset version 20', () => {
 		test.each(['x', ['x']])('input %j', input => {
-			const model = ONNXExporter.createONNXModel()
-			model.getOpsetImportList()[0].setVersion(20)
+			const model = ONNXExporter.createONNXModel({ opset: { version: 20 } })
 			gelu.export(model, { type: 'gelu', input })
 			const nodes = model.getGraph().getNodeList()
 			expect(nodes).toHaveLength(1)