Merge pull request #747 from jignparm/jignparm/fix_fusedbatchnormv3

Fix transpose_optimizer.py to handle FusedBatchNormV3 models

Author: guschmue

tf2onnx/optimizer/transpose_optimizer.py

@@ -406,16 +406,17 @@ def _add_handler(self, trans, node):
 
     def _transpose_handler(self, trans, node):
         if is_nchw_transpose(node):
-            ops = self._g.get_nodes()
-            self._g.replace_all_inputs(ops, node.output[0], trans.input[0])
-
-            shape = self._g.get_shape(node.output[0])
-            dtype = self._g.get_dtype(node.output[0])
+            for g in {self._g, node.graph}:
+                ops = g.get_nodes()
+                g.replace_all_inputs(ops, node.output[0], trans.input[0])
+
+            shape = node.graph.get_shape(node.output[0])
+            dtype = node.graph.get_dtype(node.output[0])
+            if node.output[0] in node.graph.outputs:
+                node.graph.make_node("Identity", [trans.input[0]],
+                                     outputs=node.output, shapes=[shape], dtypes=[dtype])
             self._g.remove_node(trans.name)
-            self._g.remove_node(node.name)
-            if node.output[0] in self._g.outputs:
-                self._g.make_node("Identity", [trans.input[0]],
-                                  outputs=node.output, shapes=[shape], dtypes=[dtype])
+            node.graph.remove_node(node.name)
             return True
         return False
 
@@ -459,11 +460,12 @@ def _mul_handler(self, trans, node):
         return False
 
     def _identity_handler(self, trans, node):
-        if node.output[0] in self._g.outputs:
+        if node.output[0] in node.graph.outputs:
             return False
-        ops = self._g.get_nodes()
-        self._g.replace_all_inputs(ops, node.output[0], trans.output[0])
-        self._g.remove_node(node.name)
+        for g in {self._g, node.graph}:
+            ops = g.get_nodes()
+            g.replace_all_inputs(ops, node.output[0], trans.output[0])
+        node.graph.remove_node(node.name)
         return True
 
     def _concat_handler(self, trans, node):