Add a small tool to remove cast node from onnx model (#227)

* add a decast tool in the onnxtk.

* More fixes.

* add unit test.

* update the ci build script.

* fall back the old onnxruntime for the test.

Author: Wenbing Li

.appveyor.yml

@@ -36,12 +36,10 @@ jobs:
       conda install -c conda-forge protobuf
       conda install -c conda-forge numpy
       conda install -c conda-forge cmake
-      conda install -c conda-forge openmpi
-      conda install -c conda-forge tensorflow
       pip install $(ONNX_PATH)
       pip install -r requirements.txt
       pip install -r requirements-dev.txt
-      test '$(python.version)' != '2.7' && pip install onnxruntime
+      test '$(python.version)' != '2.7' && pip install onnxruntime==0.1.4
       pip install pytest
       git clone --recursive https://github.com/cjlin1/libsvm libsvm
       cd libsvm

.azure-pipelines/linux-conda-CI.yml

@@ -1,6 +1,3 @@
-<p align="center"><img width="40%" src="../docs/ONNXMLTools_logo_main.png" /></p>
-
-
 # Introduction
 ONNXTK package enables you check and optimize the [ONNX](https://onnx.ai) models for ONNX inference engine.
 

.travis.yml

@@ -0,0 +1,64 @@
+import sys
+import onnx
+from .optimizer import LinkedNode, Solution
+
+
+def remove_cast(lnodes, op_set):
+
+    while True:
+        sln = []
+        for n_ in lnodes:
+            if n_.op_type in op_set and n_.in_single_path:
+                if n_.precedence[0].op_type == 'Cast' and n_.successor[0].op_type == 'Cast':
+                    sln.append(Solution(None, n_.precedence[0], n_.precedence[0], n_))
+                    sln.append(Solution(n_, n_.successor[0], n_.successor[0], None))
+                    break
+
+        if len(sln) == 0:
+            break
+
+        for s_ in sln:
+            lnodes = s_.apply(lnodes)
+
+    return lnodes
+
+
+def decast(origin_model, oplist):
+    """
+    remove the ONNX cast op from the specified operator.
+    :param origin_model:these
+    :param oplist:
+    :return:
+    """
+    graph = origin_model.graph
+    nodelist = list(graph.node)
+    del graph.node[:]
+
+    all_nodes = LinkedNode.build_from_onnx(nodelist,
+                                           [],
+                                           [i_.name for i_ in graph.input],
+                                           [o_.name for o_ in graph.output])
+
+    nodes = remove_cast(all_nodes, set(oplist))
+    for n_ in nodes:
+        graph.node.extend(n_.generate())
+
+    return origin_model
+
+
+def main():
+    if len(sys.argv) < 4:
+        print('decast.py model_in  model_out <op1, ...>')
+        return
+
+    input = sys.argv[1]
+    output = sys.argv[2]
+    op_list = sys.argv[3:]
+
+    oxml = onnx.load_model(input)
+    oxml = decast(oxml, op_list)
+    onnx.save_model(oxml, output)
+
+
+if __name__ == "__main__":
+    main()

onnxutils/README.md

@@ -129,6 +129,7 @@ def in_redirect(self, old_name, name):
             self.input[key] = name
 
     def out_redirect(self, old_name, name):
+        assert self.in_or_out
         if old_name in self.output:
             self.output[old_name] = name
         else:
@@ -191,10 +192,10 @@ def build_from_onnx(onnx_nodes, nchw_inputs, inputs, outputs):
                     if var_ in nchw_inputs:
                         nnode = LinkedNode(
                             helper.make_node(
-                            'Transpose',
-                            [var_],
-                            [new_output],
-                            perm=[0, 2, 3, 1]))
+                                'Transpose',
+                                [var_],
+                                [new_output],
+                                perm=[0, 2, 3, 1]))
                         var_map[new_output] = nnode
                         nnode.add_precedence(target, var_)
                         n_.in_redirect(var_, new_output)
@@ -236,6 +237,10 @@ def debug_print(node_list):
 
 
 class Solution(object):
+    """
+    Solution is the base class for solutions, and it has a basic function is to
+     delete the node range of (begin, begin_n, end_p, end), where 'begin' and 'end' are excluded.
+    """
     def __init__(self, begin, begin_n, end_p, end):
         self.begin = begin
         self.begin_n = begin_n
@@ -255,23 +260,64 @@ def is_useless_transpose(perm):
         return perm == list(six.moves.range(len(perm)))
 
     @staticmethod
-    def delete_node(node_list, begin, node, end):  # type: ([],LinkedNode, LinkedNode, LinkedNode)->[]
+    def delete_node_nto1(node_list, begin, node, end):  # type: ([],LinkedNode, LinkedNode, LinkedNode)->[]
+        """
+        delete the node which has n-input and 1-output
+        """
+        if begin is None:
+            assert node is not None
+            begin = node.precedence
+        elif not isinstance(begin, list):
+            begin = [begin]
+
         if end.in_or_out:
+            # if the end is output node, the output name will be kept to avoid the model output name updating.
+            for nb_ in begin:
+                nb_.out_redirect(node.single_input, node.single_output)
+        else:
+            for nb_ in begin:
+                target_var_name = node.single_input
+                assert target_var_name in nb_.output.values()  # since the output info never be updated, except the final.
+                end.in_redirect(node.single_output, target_var_name)
+
+        for nb_ in begin:
+            nb_.successor = [end if v_ == node else v_ for v_ in nb_.successor]
+        end.precedence = [v_ for v_ in end.precedence if v_ != node] + node.precedence
+
+        node_list.remove(node)
+        return node_list
+
+    @staticmethod
+    def delete_node_1ton(node_list, begin, node, end):  # type: ([],LinkedNode, LinkedNode, LinkedNode)->[]
+        """
+        delete the node which has 1-input and n-output
+        """
+        if end is None:
+            assert end is not None
+            end = node.successor
+        elif not isinstance(end, list):
+            end = [end]
+
+        if any(e_.in_or_out for e_ in end):
             # if the end is output node, the output name will be kept to avoid the model output name updating.
             begin.out_redirect(node.single_input, node.single_output)
         else:
-            target_var_name = node.single_input
-            assert target_var_name in begin.output.values()  # since the output info never be updated, except the final.
-            end.in_redirect(node.single_output, target_var_name)
+            for ne_ in end:
+                target_var_name = node.single_input
+                # since the output info never be updated, except the final.
+                assert target_var_name in begin.output.values()
+                ne_.in_redirect(node.single_output, target_var_name)
 
-        begin.successor = [end if v_ == node else v_ for v_ in begin.successor]
-        end.precedence = [begin if v_ == node else v_ for v_ in end.precedence]
+        begin.successor = [v_ for v_ in begin.successor if v_ != node] + node.successor
+        for ne_ in end:
+            ne_.precedence = [begin if v_ == node else v_ for v_ in ne_.precedence]
 
         node_list.remove(node)
         return node_list
 
     @staticmethod
-    def add_siso_node(node_list, begin, end, begin_output_name, node):  # type: ([], LinkedNode, LinkedNode, string, LinkedNode)->[]
+    def add_siso_node(node_list, begin, end, begin_output_name, node):
+        # type: ([], LinkedNode, LinkedNode, str, LinkedNode)->[]
         node.in_redirect(node.single_input, begin_output_name)
         end.in_redirect(begin_output_name, node.single_output)
         begin.successor[begin.successor.index(end)] = node
@@ -287,8 +333,11 @@ def apply(self, node_list):
         while node != self.end:
             assert len(node.successor) == 1
             end = node.successor[0]
-            node_list = self.delete_node(node_list, self.begin, node, end)
-            node = end
+            if self.begin:
+                node_list = self.delete_node_nto1(node_list, self.begin, node, end)
+            else:
+                node_list = self.delete_node_nto1(node_list, self.begin, node, end)
+            node = self.end if self.end is None else end
 
         return node_list
 
@@ -306,10 +355,10 @@ def apply(self, node_list):
                 #    node.reshape_input_for_broadcast(perm0)
                 node = node.successor[0]
 
-            node_list = self.delete_node(node_list, self.begin, self.begin_n, self.begin_n.successor[0])
-            node_list = self.delete_node(node_list, self.end_p.precedence[0], self.end_p, self.end)
+            node_list = self.delete_node_1ton(node_list, self.begin, self.begin_n, self.begin_n.successor[0])
+            node_list = self.delete_node_1ton(node_list, self.end_p.precedence[0], self.end_p, self.end)
         else:
-            node_list = self.delete_node(node_list, self.begin_n, self.end_p, self.end)
+            node_list = self.delete_node_1ton(node_list, self.begin_n, self.end_p, self.end)
             self.begin_n.attribute['perm'] = perm_f
         return node_list
 
@@ -346,7 +395,7 @@ def apply(self, node_list):
             FanOutSolution.number = FanOutSolution.number + 1
             node_list = Solution.add_siso_node(node_list, self.end_p, suc, list(suc.input.values())[0], nnode)
 
-        node_list = Solution.delete_node(node_list, self.begin, self.begin_n, self.end_p)
+        node_list = Solution.delete_node_1ton(node_list, self.begin, self.begin_n, self.end_p)
         return node_list
 
 
@@ -368,7 +417,7 @@ def apply(self, node_list):
         precedence_list = self.begin.precedence.copy()
         node_list = Solution.add_siso_node(node_list, self.begin, self.begin_n, list(self.begin.output.values())[0], nnode)
         for branch in precedence_list:
-            node_list = Solution.delete_node(node_list, branch.precedence[0], branch, self.begin)
+            node_list = Solution.delete_node_1ton(node_list, branch.precedence[0], branch, self.begin)
         return node_list
 
 

onnxutils/onnxtk/decast.py

@@ -0,0 +1,30 @@
+import unittest
+
+from onnx import helper
+from onnx import onnx_pb as onnx_proto
+from onnxtk.decast import decast
+
+
+class DecastTestCase(unittest.TestCase):
+
+    def test_decast(self):
+        nodes = []
+        nodes[0:] = [helper.make_node('Identity', ['input1'], ['identity1'])]
+        nodes[1:] = [helper.make_node('Cast', ['identity1'], ['cast0'], to=1)]
+        nodes[2:] = [helper.make_node('ReduceSum', ['cast0'], ['reduce0'])]
+        nodes[3:] = [helper.make_node('Cast', ['reduce0'], ['cast1'], to=6)]
+        nodes[4:] = [helper.make_node('Identity', ['cast1'], ['output0'])]
+
+        input0 = helper.make_tensor_value_info('input1', onnx_proto.TensorProto.FLOAT, [1, 1, 2, 3])
+        output0 = helper.make_tensor_value_info('output0', onnx_proto.TensorProto.FLOAT, [1, 1, 2, 3])
+
+        graph = helper.make_graph(nodes, 'test_graph', [input0], [output0])
+        model = helper.make_model(graph)
+        self.assertIsNotNone(model)
+
+        converted_model = decast(model, ['ReduceSum'])
+        self.assertTrue(len(converted_model.graph.node) == 3)
+
+
+if __name__ == '__main__':
+    unittest.main()

onnxutils/onnxtk/optimizer.py

@@ -190,13 +190,6 @@ def test_fan_in(self):
         self.assertEqual(len(new_nodes), 7)
         self.assertIsNotNone(model)
 
-    @unittest.skip('Need manually copy the real_yolov3.onnx, which is big.')
-    def test_onnx_model(self):
-        with open("./real_yolov3.onnx", 'rb') as fml:
-            oxml= onnx.load_model(fml)  # type: onnx.ModelProto
-            oxml = optimize_onnx_model(oxml)
-            onnx.save_model(oxml, 'real_yolov3_opt.onnx')
-
 
 if __name__ == '__main__':
     unittest.main()