support adding new model output, which is mentioned in #7, #8, and #13

ZhangGe6 · Sep 17, 2022 · 07ccb3f · 07ccb3f
1 parent 4cddc23
commit 07ccb3f
Show file tree

Hide file tree

Showing 13 changed files with 164 additions and 77 deletions.
diff --git a/app_desktop.py b/app_desktop.py
@@ -15,7 +15,7 @@
         - flaskwebgui github repo: https://github.com/ClimenteA/flaskwebgui
     b. add some scripts to keep server running while gui is running
         - see here: https://github.com/ClimenteA/flaskwebgui#install
-        - I added the code in the index.js
+        - I added the code in the index.js (around line 355)
     c. Then run: `python app_desktop.py`, the web browser will be automatically lauched for onnx-modifier
 
 2. How to generate excutable files:

diff --git a/docs/add_new_outputs.gif b/docs/add_new_outputs.gif
diff --git a/docs/rename_model_io.gif b/docs/rename_model_io.gif
diff --git a/docs/todo_list.md b/docs/todo_list.md
@@ -1,15 +1,17 @@
 # TODO
 
-- support desktop application.
-- support more flexible downloading schema
-    - As this [request](https://github.com/ZhangGe6/onnx-modifier/pull/5) notes, the current downloading schema prevents `onnx-modifier ` from being deployed remotely as a service.
-
-- support adding more complicated nodes (which has some simple parameters like `reshape`).
-- support adding model input/output node.
-- support combine models.
-- support user-defined input/output number when the type of node's input/output is list.
-- slim the codes.
-    - because some `.js` files (like electron.js and even python.js) in the `static` folder and `electron.html` in `templates` folder are legacy of Netron and can be further slimmed.
+- [ ] support desktop application.
+    - [x] Windows
+    - [ ] Linux
+- [ ] support more flexible downloading schema
+    - [ ] As this [request](https://github.com/ZhangGe6/onnx-modifier/pull/5) notes, the current downloading schema prevents `onnx-modifier ` from being deployed remotely as a service.
+- [ ] support adding more complicated nodes (which has some simple parameters like `reshape`).
+- [ ] support combine models.
+- [ ] support user-defined input/output number when the type of node's input/output is list.
+- [ ] slim the code.
+    - [ ] because some `.js` files (like electron.js and even python.js) in the `static` folder and `electron.html` in `templates` folder are legacy of Netron and can be further slimmed.
+- [x] support adding model input/output node.
+- [x] fix issue that "extra model inputs" emerges after deleting nodes. [issue#12](https://github.com/ZhangGe6/onnx-modifier/issues/12)
 
 
 # Some known reference issues/feature requests

diff --git a/docs/update_log.md b/docs/update_log.md
@@ -1,15 +1,23 @@
 # onnx-modifier update log
 
+## 20220813
+
+- support adding model input/output node. [issue 7](https://github.com/ZhangGe6/onnx-modifier/issues/7), [issue 8](https://github.com/ZhangGe6/onnx-modifier/issues/8), [issue 13](https://github.com/ZhangGe6/onnx-modifier/issues/13)
+
+- fix issue that "extra model inputs" emerges after deleting nodes. [issue#12](https://github.com/ZhangGe6/onnx-modifier/issues/12)
+- update windows executable file.
+
 ## 20220701
-support renaming model input/output
+
+- support renaming model input/output
 
 ## 20220621
-add Dockerfile
-    - thanks to [fengwang](https://github.com/fengwang)
+- add Dockerfile
+  - thanks to [fengwang](https://github.com/fengwang)
 
 ## 20220620
 
-add Windows executable file.
+- add Windows executable file.
 
 ## 20220612
 

diff --git a/onnx_modifier.py b/onnx_modifier.py
@@ -109,11 +109,12 @@ def remove_node_by_node_states(self, node_states):
                     # self.initializer.remove(self.initializer_name2module[init_name])
 
     def modify_node_io_name(self, node_renamed_io):
-        # print(node_renamed_io)
         for node_name in node_renamed_io.keys():
+            if node_name not in self.node_name2module.keys():
+                # custom added nodes or custom added model outputs
+                continue
             renamed_ios = node_renamed_io[node_name]
             for src_name, dst_name in renamed_ios.items():
-                # print(src_name, dst_name)
                 node = self.node_name2module[node_name]
                 if node_name in self.graph_input_names:
                     node.name = dst_name
@@ -149,16 +150,28 @@ def add_node(self, nodes_info, node_states):
 
             self.graph.node.append(node)
 
-
+    def add_outputs(self, added_outputs, node_states):
+        # https://github.com/onnx/onnx/issues/3277#issuecomment-1050600445
+        added_output_names = added_outputs.values()
+        # filter out the deleted custom-added outputs
+        value_info_protos = []
+        shape_info = onnx.shape_inference.infer_shapes(self.model_proto)
+        for value_info in shape_info.graph.value_info:
+            if value_info.name in added_output_names:
+                value_info_protos.append(value_info)
+        self.graph.output.extend(value_info_protos)
+
     def modify(self, modify_info):
         # print(modify_info['node_states'])
         # print(modify_info['node_renamed_io'])
         # print(modify_info['node_changed_attr'])
         # print(modify_info['added_node_info'])
+        # print(modify_info['added_outputs'])
         self.remove_node_by_node_states(modify_info['node_states'])
         self.modify_node_io_name(modify_info['node_renamed_io'])
         self.modify_node_attr(modify_info['node_changed_attr'])
         self.add_node(modify_info['added_node_info'], modify_info['node_states'])
+        self.add_outputs(modify_info['added_outputs'], modify_info['node_states'])
 
 
     def check_and_save_model(self, save_dir='./modified_onnx'):
@@ -191,15 +204,15 @@ def inference(self, x=None, output_names=None):
 
         # This issue may be encountered: https://github.com/microsoft/onnxruntime/issues/7506
         out = inference_session.run(None, {input_name: x})[0]
-        # print(out)
-
+        print(out.shape)
 
 if __name__ == "__main__":
     # model_path = "C:\\Users\\ZhangGe\\Desktop\\squeezenet1.0-3.onnx"
     # model_path = "C:\\Users\\ZhangGe\\Desktop\\squeezenet1.0-12-int8.onnx"
     # model_path = "C:\\Users\\ZhangGe\\Desktop\\tflite_sim.onnx"
-    # model_path = "C:\\Users\\ZhangGe\\Desktop\\modified_modified_squeezenet1.0-12.onnx"
-    model_path = "C:\\Users\\ZhangGe\\Desktop\\squeezenet1.0-3.onnx"
+    # model_path = "C:\\Users\\ZhangGe\\Desktop\\squeezenet1.0-12.onnx"
+    model_path = "C:\\Users\\ZhangGe\\Desktop\\with-added-output-modified_modified_squeezenet1.0-12.onnx"
+    # model_path = "C:\\Users\\ZhangGe\\Desktop\\squeezenet1.0-3.onnx"   // TODO: this model is not supported well , but why?
     # model_path = "C:\\Users\\ZhangGe\\Desktop\\mobilenetv2-7.onnx"
     onnx_modifier = onnxModifier.from_model_path(model_path)
 
@@ -259,7 +272,6 @@ def remove_node_by_node_states():
 
         onnx_modifier.check_and_save_model()
     # remove_node_by_node_states()
-
 
     def test_modify_node_io_name():
         node_rename_io = {'input': {'input': 'inputd'}, 'Conv_0': {'input': 'inputd'}}
@@ -285,34 +297,13 @@ def test_change_node_attr():
         onnx_modifier.check_and_save_model()
     # test_change_node_attr()
 
-
-    def debug_remove_node_by_node_states():
-        # print(len(onnx_modifier.graph.node))
-        # print(len(onnx_modifier.graph.initializer))
-
-        # print(onnx_modifier.node_name2module.keys())
-        # print(onnx_modifier.graph.node)
-        # for node in onnx_modifier.graph.node:
-        #     print(node.name)
-        #     print(node.input)
-        #     print(node.output)
-
-        # print('\noriginal input')
-        # for inp in onnx_modifier.graph.input:
-        #     print(inp.name)
-
-
-        node_states = {'data_0': 'Exist', 'Conv0': 'Exist', 'Relu1': 'Exist', 'MaxPool2': 'Exist', 'Conv3': 'Exist', 'Relu4': 'Exist', 'Conv5': 'Exist', 'Relu6': 'Exist', 'Conv7': 'Exist', 'Relu8': 'Exist', 'Concat9': 'Exist', 'Conv10': 'Exist', 'Relu11': 'Exist', 'Conv12': 'Exist', 'Relu13': 'Exist', 'Conv14': 'Exist', 'Relu15': 'Exist', 'Concat16': 'Exist', 'MaxPool17': 'Exist', 'Conv18': 'Exist', 'Relu19': 'Exist', 'Conv20': 'Exist', 'Relu21': 'Exist', 'Conv22': 'Exist', 'Relu23': 'Exist', 'Concat24': 'Exist', 'Conv25': 'Exist', 'Relu26': 'Exist', 'Conv27': 'Exist', 'Relu28': 'Exist', 'Conv29': 'Exist', 'Relu30': 'Exist', 'Concat31': 'Exist', 'MaxPool32': 'Exist', 'Conv33': 'Exist', 'Relu34': 'Exist', 'Conv35': 'Exist', 'Relu36': 'Exist', 'Conv37': 'Exist', 'Relu38': 'Exist', 'Concat39': 'Exist', 'Conv40': 'Exist', 'Relu41': 'Exist', 'Conv42': 'Exist', 'Relu43': 'Exist', 'Conv44': 'Exist', 'Relu45': 'Exist', 'Concat46': 'Exist', 'Conv47': 'Exist', 'Relu48': 'Exist', 'Conv49': 'Exist', 'Relu50': 'Exist', 'Conv51': 'Exist', 'Relu52': 'Exist', 'Concat53': 'Exist', 'Conv54': 'Exist', 'Relu55': 'Deleted', 'Conv56': 'Deleted', 'Relu57': 'Deleted', 'Conv58': 'Deleted', 'Relu59': 'Deleted', 'Concat60': 'Deleted', 'Dropout61': 'Deleted', 'Conv62': 'Deleted', 'Relu63': 'Deleted', 'GlobalAveragePool64': 'Deleted', 'Softmax65': 'Deleted', 'out_softmaxout_1': 'Deleted'}
-        # print('\graph  input')
-        # for inp in onnx_modifier.graph.input:
-        #     print(inp.name)
-        onnx_modifier.remove_node_by_node_states(node_states)
-
-        print('\nleft input')
-        for inp in onnx_modifier.graph.input:
-            print(inp.name)
-
-        onnx_modifier.check_and_save_model()
-    debug_remove_node_by_node_states()
+    def test_inference():
+        onnx_modifier.inference()
+    # test_inference()
 
-
+    def test_add_output():
+        # print(onnx_modifier.graph.output)
+        onnx_modifier.add_outputs(['fire2/squeeze1x1_1'])
+        # print(onnx_modifier.graph.output)
+        onnx_modifier.check_and_save_model()
+    # test_add_output()
diff --git a/readme.md b/readme.md
@@ -20,6 +20,7 @@ Currently, the following editing operations are supported:
   - Recover a deleted node.
 - [x] Rename the node inputs/outputs
 - [x] Rename the model inputs/outputs
+- [x] Add new model outputs
 - [x] Edit the attribute of nodes
 - [x] Add new nodes (experimental)
 
@@ -52,8 +53,8 @@ Click the url in the output info generated by flask (`http://127.0.0.1:5000/` fo
 ## launch from executable file
 <details>
   <summary>Click to expand</summary>
-  
-- Windows: Download [onnx-modifier.exe (27.6MB)](https://drive.google.com/file/d/1y7mYlvF0G5iiNDgOFh1ESXlTs8I2ipVQ/view?usp=sharing), double-click it and enjoy.
+
+- Windows: Download [onnx-modifier.exe (27.6MB)](https://drive.google.com/file/d/1LRXgZauQ5BUENe_PvilRW8WvSO-4Jr9j/view?usp=sharing), double-click it and enjoy.
   - Edge browser is used for runtime environment by default.
 
 > I recorded how I made the the executable file in `app_desktop.py`. The executable file for other platforms are left for future work.
@@ -139,6 +140,16 @@ The process is shown in the following figure:
 ## Rename the model inputs/outputs
 Click the model input/output node, type a new name in the sidebar, then we are done.
 
+![rename_model_io](./docs/rename_model_io.gif)
+
+## Add new model outputs
+
+Sometimes we want to add/extract the output of a certain node as model output. For example, we want to add a new model output after the old one was deleted, or extract intermediate layer output for fine-grained analysis. In `onnx-modifier`, we can achieve this by simply clicking the `Add Output` button in the sidebar of the corresponding node. Then we can get a new model output node following the corresponding node. Its name is the same as the output of the corresponding node.  
+
+In the following example, we add 2 new model outputs, which are the outputs of the 1st `Conv` node and 2nd `Conv` node, respectively.
+
+![add_new_outputs](./docs/add_new_outputs.gif)
+
 ## Edit the attribute of nodes
 
 Change the original attribute to a new value, then we are done.

diff --git a/readme_zh-CN.md b/readme_zh-CN.md
@@ -20,6 +20,8 @@
   - 删除一个节点及所有以该节点为根节点的子节点
   - 恢复误删的节点
 - [x] 修改节点输入输出名
+- [x] 修改模型输入输出名
+- [x] 增加模型输出节点
 - [x] 编辑节点属性值
 - [x] 增加新节点（beta）
 
@@ -47,7 +49,7 @@
   ```
 
 ## 从可执行文件启动
-- Windows: 下载可执行文件[onnx-modifier.exe (27.6MB)](https://drive.google.com/file/d/1y7mYlvF0G5iiNDgOFh1ESXlTs8I2ipVQ/view?usp=sharing)，双击即可启动。
+- Windows: 下载可执行文件[onnx-modifier.exe (27.6MB)](https://drive.google.com/file/d/1LRXgZauQ5BUENe_PvilRW8WvSO-4Jr9j/view?usp=sharing)，双击即可启动。
   - 默认使用Edge浏览器作为运行环境。
 > 生成可执行文件的步骤记录在`app_desktop.py`文件中。未来会为其他平台生成可执行文件。
 
@@ -102,6 +104,22 @@
 
 <img src="./docs/rename_io.gif" style="zoom:75%;" />
 
+## 修改模型输入输出名称
+
+点击模型输入或输出节点，在弹出的侧边栏中，为模型输入输出键入新的名称即可。
+
+![rename_model_io](./docs/rename_model_io.gif)
+
+## 增加模型输出节点
+
+有时候我们需要增加/抽取某个特定节点的输出作为整个模型的输出。比如之前的模型输出节点在编辑过程中被删除了，需要增加新的，或者有时候我们需要抽取一些中间层特征输出做更细致的分析。
+
+通过`onnx-modifier`，我们只需要在对应节点的侧边栏中，点击`Add Output`按钮即可在该节点后部增加一个模型输出节点，其名称与原节点的输出名相同。
+
+如下图，我们增加了两个模型输出节点，分别为第一个卷积层的输出和第二个卷积层的输出。
+
+![add_new_outputs](./docs/add_new_outputs.gif)
+
 ## 编辑节点属性值
 
 在节点侧边栏对应的属性值输入框中，键入新的属性值即可。

diff --git a/static/index.js b/static/index.js
@@ -229,10 +229,10 @@ host.BrowserHost = class {
                     'node_states' : this.mapToObjectRec(this._view._graph._modelNodeName2State),
                     'node_renamed_io' : this.mapToObjectRec(this._view._graph._renameMap),
                     'node_changed_attr' : this.mapToObjectRec(this._view._graph._changedAttributes),
-                    'added_node_info' : this.mapToObjectRec(this.parseLightNodeInfo2Map(this._view._graph._addedNode))
-                }
-
-                )
+                    'added_node_info' : this.mapToObjectRec(this.parseLightNodeInfo2Map(this._view._graph._addedNode)),
+                    'added_outputs' : this.arrayToObject(this.process_added_outputs(this._view._graph._addedOutputs, 
+                                                                this._view._graph._renameMap, this._view._graph._modelNodeName2State))
+                })
             }).then(function (response) {
                 return response.text();
             }).then(function (text) {
@@ -264,9 +264,6 @@ host.BrowserHost = class {
             this._view._updateGraph();
         })
 
-
-
-
         this.document.getElementById('version').innerText = this.version;
 
         if (this._meta.file) {
@@ -677,6 +674,33 @@ host.BrowserHost = class {
         }
         return lo
     }
+
+    // this function does 2 things:
+    // 1. rename the addedOutputs with their new names using renameMap. Because addedOutputs are stored in lists,
+    //    it may be not easy to rename them while editing. (Of course there may be a better way to do this)
+    // 2. filter out the custom output which is added, but deleted later 
+    process_added_outputs(addedOutputs, renameMap, modelNodeName2State) {
+        var processed = []
+        for (let i = 0; i < addedOutputs.length; ++i) {
+            if (modelNodeName2State.get("out_" + addedOutputs[i]) == "Exist") {
+                processed.push(addedOutputs[i]);
+            }
+        }
+        for (let i = 0; i < processed.length; ++i) {
+            if (renameMap.get("out_" + processed[i])) {
+                processed[i] = renameMap.get("out_" + processed[i]).get(processed[i]);
+            }
+        }
+        return processed;
+    }
+
+    // https://stackoverflow.com/a/4215753/10096987
+    arrayToObject(arr) {
+        var rv = {};
+        for (var i = 0; i < arr.length; ++i)
+          if (arr[i] !== undefined) rv[i] = arr[i];
+        return rv;
+    }
 
     // convert view.LightNodeInfo to Map object for easier transmission to Python backend
     parseLightNodeInfo2Map(nodes_info) {

diff --git a/static/onnx.js b/static/onnx.js
@@ -438,6 +438,7 @@ onnx.Graph = class {
 
         this._custom_add_node_io_idx = 0
         this._custom_added_node = []
+        this._custom_added_outputs = []
 
         // model parameter assignment here!
         // console.log(graph)
@@ -504,7 +505,8 @@ onnx.Graph = class {
     }
 
     get outputs() {
-        return this._outputs;
+        // return this._outputs;
+        return this._outputs.concat(this._custom_added_outputs);
     }
 
     get nodes() {
@@ -632,6 +634,15 @@ onnx.Graph = class {
         return custom_add_node;
     }
 
+    reset_custom_added_outputs() {
+        this._custom_added_outputs = [];
+    }
+
+    add_output(name) {
+        const argument = this._context.argument(name);
+        this._custom_added_outputs.push(new onnx.Parameter(name, [ argument ]));
+    }
+
 };
 
 onnx.Parameter = class {

diff --git a/static/view-grapher.js b/static/view-grapher.js
@@ -24,6 +24,7 @@ grapher.Graph = class {
         this._changedAttributes = new Map();
 
         this._addedNode = new Map();
+        this._addedOutputs = [];
     }
 
     get options() {

diff --git a/static/view-sidebar.js b/static/view-sidebar.js
@@ -209,6 +209,9 @@ sidebar.NodeSidebar = class {
         this._addButton('Recover Node');
         this.add_separator(this._elements, 'sidebar-view-separator')
         this._addButton('Enter');
+
+        this._addHeader('Output adding helper');
+        this._addButton('Add Output');
 
         // deprecated
         // this.add_separator(this._elements, 'sidebar-view-separator');
@@ -272,8 +275,6 @@ sidebar.NodeSidebar = class {
                 }
             }
         }
-
-
     }
 
     render() {
@@ -356,7 +357,11 @@ sidebar.NodeSidebar = class {
                 this._host._view._updateGraph()
             });
         }
-
+        if (title === 'Add Output') {
+            buttonElement.addEventListener('click', () => {
+                this._host._view._graph.add_output(this._modelNodeName)
+            });   
+        }
     }
 
     // deprecated