ONNX loader shared edge objects (#71)

lutzroeder · Jun 17, 2018 · 3556f09 · 3556f09
1 parent c240370
commit 3556f09
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 42 deletions.
diff --git a/src/onnx-model.js b/src/onnx-model.js
@@ -87,7 +87,7 @@ class OnnxModel {
         if (this._modelVersion) {
             results.push({ name: 'Version', value: this._modelVersion });
         }
-        if (this._model.docString) {
+        if (this._docString) {
             results.push({ name: 'Description', value: this._docString });
         }
         var metadata = {};
@@ -140,8 +140,8 @@ class OnnxGraph {
         this._nodes = [];
 
         if (graph) {
-            var initializerMap = [];
-            var valueInfoMap = [];
+            this._initializerMap = {};
+            this._connectionMap = {};
 
             this._name = graph.name || ('(' + index.toString() + ')');
             this._description = graph.docString || '';
@@ -160,7 +160,7 @@ class OnnxGraph {
                     if (outputCountMap[name] == 1) {
                         var attribute = node.attribute.find((attribute) => { return attribute.name == 'value' && attribute.t; }); 
                         if (attribute) {
-                            initializerMap[name] = new OnnxTensor(attribute.t, name, 'Constant');
+                            this._initializerMap[name] = new OnnxTensor(attribute.t, name, 'Constant');
                             initializerNode = true;
                         }
                     }
@@ -171,71 +171,51 @@ class OnnxGraph {
             });
 
             graph.initializer.forEach((tensor) => {
-                initializerMap[tensor.name] = new OnnxTensor(tensor, tensor.name, 'Initializer');
+                this._initializerMap[tensor.name] = new OnnxTensor(tensor, tensor.name, 'Initializer');
             });
             graph.valueInfo.forEach((valueInfo) => {
-                valueInfoMap[valueInfo.name] = valueInfo;
+                this._connection(valueInfo.name, valueInfo.type, valueInfo.docString);
             });
 
             this._inputs = [];
             graph.input.forEach((valueInfo) => {
-                if (!initializerMap[valueInfo.name]) {
-                    this._inputs.push({
-                        id: valueInfo.name,
-                        name: valueInfo.name,
-                        description: valueInfo.docString,
-                        type: OnnxTensor.formatType(valueInfo.type)
-                    });
-                    valueInfoMap[valueInfo.name] = valueInfo;
+                if (!this._initializerMap[valueInfo.name]) {
+                    var connection = this._connection(valueInfo.name, valueInfo.type, valueInfo.docString);
+                    connection.name = valueInfo.name;
+                    this._inputs.push(connection);
                 }
             });
-
-            this._outputs = graph.output.map((valueInfo) => {
-                valueInfoMap[valueInfo.name] = valueInfo;
-                return {
-                    id: valueInfo.name,
-                    name: valueInfo.name,
-                    description: valueInfo.docString,
-                    type: OnnxTensor.formatType(valueInfo.type)
-                };
+            this._outputs = [];
+            graph.output.map((valueInfo) => {
+                var connection = this._connection(valueInfo.name, valueInfo.type, valueInfo.docString);
+                connection.name = valueInfo.name;
+                this._outputs.push(connection);
             });
 
             nodes.forEach((node) => {
                 var inputs = [];
                 if (node.input) {
                     inputs = this._metadata.getInputs(node.opType, node.input);
                     inputs.forEach((input) => {
-                        input.connections.forEach((connection) => {
-                            var initializer = initializerMap[connection.id];
-                            if (initializer) {
-                                connection.initializer = initializer;
-                                connection.type = initializer.type;
-                            }
-                            else {
-                                var valueInfo = valueInfoMap[connection.id];
-                                if (valueInfo) {
-                                    connection.type = OnnxTensor.formatType(valueInfo.type);
-                                    connection.description = valueInfo.docString;
-                                }
-                            }
+                        input.connections = input.connections.map((connection) => {
+                            return this._connection(connection.id);
                         });
                     });          
                 }
                 var outputs = [];
                 if (node.output) {
                     outputs = this._metadata.getOutputs(node.opType, node.output);
                     outputs.forEach((output) => {
-                        output.connections.forEach((connection) => {
-                            var valueInfo = valueInfoMap[connection.id];
-                            if (valueInfo) {
-                                connection.type = OnnxTensor.formatType(valueInfo.type);
-                                connection.description = valueInfo.docString;
-                            }
+                        output.connections = output.connections.map((connection) => {
+                            return this._connection(connection.id);
                         });
                     });
                 }
                 this._nodes.push(new OnnxNode(this, node.opType, node.domain, node.name, node.docString, node.attribute, inputs, outputs));
             });
+
+            delete this._initializerMap;
+            delete this._connectionMap;
         }
     }
 
@@ -266,6 +246,27 @@ class OnnxGraph {
     get metadata() {
         return this._metadata;
     }
+
+    _connection(name, type, docString) {
+        var connection = this._connectionMap[name];
+        if (!connection) {
+            connection = {};
+            connection.id = name;
+            var initializer = this._initializerMap[name];
+            if (initializer) {
+                connection.initializer = initializer;
+                connection.type = initializer.type;
+            }
+            if (type) {
+                connection.type = OnnxTensor.formatType(type);
+            }
+            if (docString) {
+                connection.description = docString;
+            }
+            this._connectionMap[name] = connection;
+        }
+        return connection;
+    }
 }
 
 class OnnxNode {

diff --git a/tools/onnx-converter b/tools/onnx-converter
@@ -23,6 +23,9 @@ fi
 export PYTHONUSERBASE=${build}/third_party/pypi/${identifier}
 export PATH=$PATH:${PYTHONUSERBASE}/bin
 rm -rf ${PYTHONUSERBASE}
+pip install --quiet --user coremltools
+pip install --quiet --user tensorflow
+pip install --quiet --user keras
 pip install --quiet --user ${third_party}/${identifier}
 
 python ${tools}/onnx-converter.py $@