Merge pull request #522 from nbcsm/yolo

fix conv_convert_inputs bug

Author: guschmue

tests/run_pretrained_models.py

@@ -251,8 +251,8 @@ def run_test(self, name, backend="caffe2", onnx_file=None, opset=None, extra_ops
                 # convert model to onnx
                 onnx_graph = self.to_onnx(sess.graph, opset=opset, extra_opset=extra_opset,
                                           shape_override=shape_override, input_names=inputs.keys())
+                onnx_graph = optimizer.optimize_graph(onnx_graph)
                 model_proto = onnx_graph.make_model("converted from tf2onnx")
-                model_proto = optimizer.optimize_graph(onnx_graph).make_model("optimized")
                 logger.info("To_ONNX, OK")
                 if onnx_file:
                     self.create_onnx_file(name, model_proto, inputs, onnx_file)

tf2onnx/onnx_opset/nn.py

@@ -55,11 +55,10 @@ def conv_convert_inputs(ctx, node, with_kernel=False, new_kernel_shape=None,
         # transpose input if needed, no need to record shapes on input
         for idx in input_indices:
             parent = node.inputs[idx]
-            if node.inputs[idx].is_const():
-                # if input is a constant, transpose that one
-                if not parent.data_format:
-                    val = parent.get_tensor_value(as_list=False)
-                    parent.set_tensor_value(val.transpose(constants.NHWC_TO_NCHW))
+            if node.inputs[idx].is_const() and len(ctx.find_output_consumers(node.input[1])) == 1:
+                # if input is a constant, transpose that one if we are the only consumer
+                val = parent.get_tensor_value(as_list=False)
+                parent.set_tensor_value(val.transpose(constants.NHWC_TO_NCHW))
             else:
                 # if input comes from a op, insert transpose op
                 input_name = node.input[idx]
@@ -70,33 +69,27 @@ def conv_convert_inputs(ctx, node, with_kernel=False, new_kernel_shape=None,
                 if shape is not None:
                     new_shape = spatial_map(shape, constants.NHWC_TO_NCHW)
                     ctx.set_shape(transpose.output[0], new_shape)
-            parent.data_format = "NCHW"
 
     # kernel must to be transposed
     if with_kernel:
         parent = node.inputs[1]
         need_transpose = True
         if node.inputs[1].is_const():
             # kernel is const - transpose the const if we are the only consumer of const
-            # TODO: maybe we should make a copy of the const, or look at the other consumers
-            # if they'd want a transose as well.
             consumers = ctx.find_output_consumers(node.input[1])
             if len(consumers) == 1:
                 val = parent.get_tensor_value(as_list=False)
                 val = val.transpose(constants.HWCN_TO_NCHW)
                 parent.set_tensor_value(val)
-                parent.data_format = "NCHW"
                 need_transpose = False
 
         if need_transpose:
             input_name = node.input[1]
             transpose = ctx.insert_new_node_on_input(node, "Transpose", input_name)
             transpose.set_attr("perm", constants.HWCN_TO_NCHW)
             transpose.skip_conversion = True
-            ctx.copy_shape(input_name, transpose.output[0])
             new_shape = spatial_map(ctx.get_shape(input_name), constants.HWCN_TO_NCHW)
             ctx.set_shape(transpose.output[0], new_shape)
-            parent.data_format = "NCHW"
 
         # some onnx conv ops require the reshape the kernel (ie. depthwise_conv2d)
         if new_kernel_shape:
@@ -129,7 +122,7 @@ def conv_convert_inputs(ctx, node, with_kernel=False, new_kernel_shape=None,
             ctx.set_shape(transpose.output[0], output_shape)
             # Transpose TF NHWC shape back to NCHW shape for current ONNX conv node output
             ctx.set_shape(output_name, spatial_map(output_shape, constants.NHWC_TO_NCHW))
-            node.data_format = "NCHW"
+        node.data_format = "NCHW"
 
 
 def add_padding(ctx, node, kernel_shape, strides, dilations=None, spatial=2):

tf2onnx/optimizer/__init__.py

@@ -50,7 +50,7 @@ def optimize_graph(graph):
     diff = copy.deepcopy(after)
     diff.subtract(before)
     diff = ["{} {} ({}->{})".format(k, str(v) if v < 0 else '+' + str(v), before.get(k, 0), after.get(k, 0))
-            for k, v in diff.most_common() if v != 0]
+            for k, v in sorted(diff.items()) if v != 0]
     logger.info("After optimization: %s", ', '.join(diff) if diff else "no change")
 
     return graph

tf2onnx/optimizer/const_fold_optimizer.py

@@ -69,7 +69,7 @@ def _fold_node(self, node, graph):
                 const_outputs = process_func(node, graph)
                 self._replace_node_with_const(node, graph, const_outputs)
                 return True
-        self.logger.debug("need to add function to fold op %s whose op_type is %s", node.name, node.type)
+            self.logger.debug("need to add function to fold op %s whose op_type is %s", node.name, node.type)
         return False
 
     @staticmethod

tf2onnx/optimizer/optimizer_base.py

@@ -63,5 +63,5 @@ def _print_stat_diff(self, before, after):
         diff = copy.deepcopy(after)
         diff.subtract(before)
         diff = ["{} {} ({}->{})".format(k, str(v) if v < 0 else '+' + str(v), before.get(k, 0), after.get(k, 0))
-                for k, v in diff.most_common() if v != 0]
+                for k, v in sorted(diff.items()) if v != 0]
         self.logger.verbose(', '.join(diff) if diff else "no change")

tf2onnx/shape_inference.py

@@ -24,7 +24,9 @@
 def infer_shape(tf_graph, shape_override):
     """Infer shape for TF graph with shape_override set first."""
     if shape_override:
+        logger.info("Apply shape override:")
         for name, shape in shape_override.items():
+            logger.info("\tSet %s shape to %s", name, shape)
             tf_graph.get_tensor_by_name(name).set_shape(shape)
         tf_graph = reload_tf_graph(tf_graph)