Fix transpose optimizer to convert simple FusedBatchNormV3 models correctly. The Identity and Transpose handlers require to update both the subgraph as well as parent graph, and also cases where an op in parent graph is consumed in subgraph.

Author: jignparm

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):