refactor

Author: zhijxu-MS

tests/common.py

@@ -4,14 +4,14 @@
 """ test common utilities."""
 
 import argparse
-import numpy as np
 import os
 import sys
 import unittest
 from collections import defaultdict
 
 from distutils.version import LooseVersion
 from parameterized import parameterized
+import numpy as np
 from tf2onnx import constants, logging, utils
 
 __all__ = [
@@ -281,11 +281,8 @@ def check_onnxruntime_incompatibility(op):
 def validate_const_node(node, expected_val):
     if node.is_const():
         node_val = node.get_tensor_value()
-        if (isinstance(expected_val, list) and isinstance(expected_val[0], float)) \
-           or isinstance(expected_val, float):
-            np.testing.assert_allclose(expected_val, node_val)
-            return True
-        return node_val == expected_val
+        np.testing.assert_allclose(expected_val, node_val)
+        return True
     return False
 
 

tf2onnx/onnx_opset/nn.py

@@ -150,8 +150,7 @@ def add_padding(ctx, node, kernel_shape, strides, dilations=None, spatial=2):
                 input_shape = spatial_map(input_shape, constants.NHWC_TO_NCHW)
                 output_shape = spatial_map(output_shape, constants.NHWC_TO_NCHW)
             # calculate pads
-            if any(input_shape[i + 2] == -1 or output_shape[i + 2] == -1 for i in range(spatial)) \
-               or any(output_shape[i + 2] == -1 for i in range(spatial)):
+            if any(input_shape[i + 2] == -1 or output_shape[i + 2] == -1 for i in range(spatial)):
                 logger.debug("node %s has unknown dim %d for pads calculation, fallback to auto_pad",
                              node.name, input_shape)
                 node.set_attr("auto_pad", "SAME_UPPER")

tf2onnx/shape_inference.py

@@ -121,7 +121,8 @@ def infer_shape_for_node(g, node):
         # https://www.tensorflow.org/api_docs/python/tf/gather
         shape_params = g.get_shape(node.input[0])
         shape_indices = g.get_shape(node.input[1])
-        # in lower tf version, gather only has 2 inputs
+        # gather can only have 2 inputs
+        # https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/gather.html
         if len(node.input) == 3:
             axis = node.input[2].get_tensor_value()
         else:
@@ -218,7 +219,7 @@ def infer_output_shapes_with_partial_inputs(g, node):
         data_inputs = node.input[:-1]
         input_shapes = [g.get_shape(node) for node in data_inputs]
         input_shapes = [shape for shape in input_shapes if shape is not None]
-        if len(input_shapes) == 0:
+        if not input_shapes:
             logger.debug("all input shapes of concat node %s are None, can't infer its output shape", node.name)
             return False