preserve output order when converting batchnorm+scale layers (#571)

microsoft · Feb 13, 2019 · 18812bf · 18812bf
1 parent ec332ff
commit 18812bf
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/mmdnn/conversion/caffe/graph.py b/mmdnn/conversion/caffe/graph.py
@@ -147,9 +147,10 @@ def __init__(self, name, kind, layer=None):
         self.output_shape = None
         self.metadata = {}
 
-    def add_parent(self, parent_node, from_output):
+    def add_parent(self, parent_node, from_output, index=None):
         assert parent_node not in self.parents
-        self.parents.append((parent_node, from_output))
+        index = len(self.parents) if index is None else index
+        self.parents.insert(index, (parent_node, from_output))
         if self not in parent_node.children:
             parent_node.children.append(self)
 

diff --git a/mmdnn/conversion/caffe/transformer.py b/mmdnn/conversion/caffe/transformer.py
@@ -176,8 +176,9 @@ def __call__(self, graph):
                 continue
             # Rewrite the fused node's children to its parent.
             for child in node.children:
-                child.parents = [(input, idx) for input, idx in child.parents if input != node]
-                child.add_parent(parent, from_output)
+                index = [n for n, (input, idx) in enumerate(child.parents) if input == node][0]
+                child.parents.pop(index)
+                child.add_parent(parent, from_output, index)
             # Disconnect the fused node from the graph.
             parent.children.remove(node)
             fused_nodes.append(node)