support, test for fusedbatchnorm

Author: guschmue

tests/run_pretrained_models.py

@@ -166,7 +166,7 @@ def run_tensorflow(self, sess, inputs):
     @staticmethod
     def to_onnx(tf_graph, opset=None):
         """Convert graph to tensorflow."""
-        return process_tf_graph(tf_graph, opset=opset)
+        return process_tf_graph(tf_graph, continue_on_error=False, opset=opset)
 
     def run_caffe2(self, name, onnx_graph, inputs):
         """Run test again caffe2 backend."""
@@ -271,6 +271,7 @@ def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=No
         graph_def = graph_pb2.GraphDef()
         with open(model_path, "rb") as f:
             graph_def.ParseFromString(f.read())
+        graph_def = tf2onnx.tfonnx.tf_optimize(None, inputs, self.output_names, graph_def)
 
         g = tf.import_graph_def(graph_def, name='')
         with tf.Session(graph=g) as sess:

tests/run_pretrained_models.yaml

@@ -25,6 +25,16 @@ benchtf-conv:
   outputs:
     - output:0
 
+benchtf-convbn:
+  disabled: true
+  # fails with: expects to be colocated with unknown node 'batch_normalization_1/gamma
+  model: tests/models/convbn-layers/frozen.pb
+  input_get: get_ramp
+  inputs:
+    "X:0": [1, 784]
+  outputs:
+    - output:0
+
 benchtf-ae0:
   model: tests/models/ae0/frozen.pb
   input_get: get_ramp

tests/test_backend.py

@@ -95,6 +95,7 @@ def run_onnxmsrtnext(onnx_graph, inputs, output_names, test_name):
         """Run test against msrt-next backend."""
         import lotus
         model_path = os.path.join(TMPPATH, test_name + ".pb")
+        # print(model_path)
         with open(model_path, "wb") as f:
             f.write(onnx_graph.SerializeToString())
         m = lotus.InferenceSession(model_path)
@@ -801,7 +802,6 @@ def test_cast(self):
         self.assertAllClose(expected, actual)
 
     def test_onehot0(self):
-        # no such op in onnx
         x_val = np.array([0, 1, 2], dtype=np.int32)
         depth = 5
         for axis in [-1, 0, 1]:
@@ -814,7 +814,7 @@ def test_onehot0(self):
 
     @unittest.skip
     def test_onehot1(self):
-        # no such op in onnx
+        # only rank 1 is currently implemented
         x_val = np.array([[0, 2], [1, -1]], dtype=np.int32)
         depth = 3
         x = tf.placeholder(tf.int32, x_val.shape, name=_TFINPUT)
@@ -824,7 +824,6 @@ def test_onehot1(self):
         self.assertAllClose(expected, actual)
 
     def test_onehot2(self):
-        # no such op in onnx
         x_val = np.array([0, 1, 2, 1, 2, 0, 1, 2, 1, 2], dtype=np.int32)
         depth = 20
         x = tf.placeholder(tf.int32, x_val.shape, name=_TFINPUT)
@@ -924,6 +923,60 @@ def test_strided_slice2(self):
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
 
+    @unittest.skip
+    def test_strided_slice3(self):
+        x_val = np.arange(3*2*3).astype("float32").reshape(3, 2, 3)
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        x_ = x[1:]
+        output = tf.identity(x_, name=_TFOUTPUT)
+        actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
+        self.assertAllClose(expected, actual)
+
+    @unittest.skip
+    def test_strided_slice4(self):
+        x_val = np.arange(3*2*3).astype("float32").reshape(3, 2, 3)
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        x_ = x[:2]
+        output = tf.identity(x_, name=_TFOUTPUT)
+        actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
+        self.assertAllClose(expected, actual)
+
+    @unittest.skip
+    def test_strided_slice5(self):
+        x_val = np.arange(3*2*3).astype("float32").reshape(3, 2, 3)
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        x_ = x[:2, 0:1, 1:]
+        output = tf.identity(x_, name=_TFOUTPUT)
+        actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
+        self.assertAllClose(expected, actual)
+
+    @unittest.skipIf(BACKEND in ["caffe2", "onnxmsrt"], "fails with schema error")
+    def test_batchnorm(self):
+        x_shape = [1, 28, 28, 2]
+        x_dtype = np.float32
+        scale_dtype = np.float32
+        scale_shape =  [2]
+        # only nhwc is support on cpu for tensorflow
+        data_format = "NHWC"
+        x_val = np.random.random_sample(x_shape).astype(x_dtype)
+        scale_val = np.random.random_sample(scale_shape).astype(scale_dtype)
+        offset_val = np.random.random_sample(scale_shape).astype(scale_dtype)
+        mean_val = np.random.random_sample(scale_shape).astype(scale_dtype)
+        var_val = np.random.random_sample(scale_shape).astype(scale_dtype)
+
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        scale = tf.constant(scale_val, name='scale')
+        offset = tf.constant(offset_val, name='offset')
+        mean = tf.constant(mean_val, name='mean')
+        var = tf.constant(var_val, name='variance')
+        epsilon = 0.001
+        y, _, _ = tf.nn.fused_batch_norm(
+            x, scale, offset, mean=mean, variance=var,
+            epsilon=epsilon, data_format=data_format, is_training=False)
+        output = tf.identity(y, name=_TFOUTPUT)
+        actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
+        self.assertAllClose(expected, actual, rtol=1e-04)
+
     @unittest.skipIf(BACKEND in ["caffe2", "onnxmsrt"], "not correctly supported")
     def test_resize_nearest_neighbor(self):
         x_shape = [1, 15, 20, 2]
@@ -964,10 +1017,13 @@ def test_fill(self):
     parser.add_argument('--backend', default='caffe2',
                         choices=["caffe2", "onnxmsrt", "onnxmsrtnext", "onnx-tensorflow"],
                         help="backend to test against")
+    parser.add_argument('--opset', default=OPSET,
+                        help="opset to test against")
     parser.add_argument('unittest_args', nargs='*')
 
     args = parser.parse_args()
     BACKEND = args.backend
+    OPSET = args.opset
     # Now set the sys.argv to the unittest_args (leaving sys.argv[0] alone)
     sys.argv[1:] = args.unittest_args
     unittest.main()

tf2onnx/tfonnx.py

@@ -305,7 +305,8 @@ def reshape_op5(ctx, node, name, args):
 NCHW_TO_HWCN = [2, 3, 1, 0]
 
 
-def conv_convert_inputs(ctx, node, with_kernel=False, new_kernel_shape=None):
+def conv_convert_inputs(ctx, node, with_kernel=False, new_kernel_shape=None,
+                        input_indices=None, output_indices=None):
     """Convert input and kernel from tensorflow to onnx. This maybe require to
         to insert transpose ops for input, kernel and output unless they are constants
         and we can transpose the constant.
@@ -324,25 +325,32 @@ def calc_shape(a, b):
             return [a[b[i]] for i in b]
         return None
 
+    if input_indices is None:
+        input_indices = [0]
+    if output_indices is None:
+        output_indices = [0]
+
     nodes = []
 
     if node.is_nhwc():
         # transpose input if needed, no need to record shapes on input
-        if node.inputs[0].is_const():
-            # if input is a constant, transpose that one
-            parent = node.inputs[0]
-            if not parent.data_format:
-                val = parent.get_tensor_value()
-                parent.set_tensor_value(val.transpose(NHWC_TO_NCHW))
-                parent.data_format = "NCHW"
-        else:
-            # if input comes from a op, insert transpose op
-            input_name = node.input[0]
-            transpose = ctx.insert_new_node_on_input(node, "Transpose", input_name)
-            transpose.set_attr("perm", NHWC_TO_NCHW)
-            transpose.inserted_nchw = True
-            ctx.set_shape(transpose.output[0], calc_shape(ctx.get_shape(input_name), NHWC_TO_NCHW))
-            nodes.append(transpose)
+        for idx in input_indices:
+            if node.inputs[idx].is_const():
+                # if input is a constant, transpose that one
+                parent = node.inputs[idx]
+                if not parent.data_format:
+                    val = parent.get_tensor_value()
+                    parent.set_tensor_value(val.transpose(NHWC_TO_NCHW))
+                    parent.data_format = "NCHW"
+            else:
+                # if input comes from a op, insert transpose op
+                input_name = node.input[idx]
+                transpose = ctx.insert_new_node_on_input(node, "Transpose", input_name)
+                transpose.set_attr("perm", NHWC_TO_NCHW)
+                transpose.inserted_nchw = True
+                if idx == 0:
+                    ctx.set_shape(transpose.output[0], calc_shape(ctx.get_shape(input_name), NHWC_TO_NCHW))
+                nodes.append(transpose)
 
     # kernel mist to be transposed
     if with_kernel:
@@ -388,12 +396,13 @@ def calc_shape(a, b):
     # transpose outputs if needed
     if node.is_nhwc():
         # TODO: what if len(output) > 0 ?
-        for i, output_name in enumerate(node.output):
+        for idx in output_indices:
+            output_name = node.output[idx]
             op_name = utils.make_name(node.name)
             transpose = ctx.insert_new_node_on_output("Transpose", output_name, name=op_name)
             transpose.set_attr("perm", NCHW_TO_NHWC)
             transpose.inserted_nchw = True
-            ctx.set_shape(transpose.output[0], calc_shape(ctx.get_shape(node.output[0]), NCHW_TO_NHWC))
+            ctx.set_shape(transpose.output[0], calc_shape(ctx.get_shape(node.output[idx]), NCHW_TO_NHWC))
             nodes.append(transpose)
     return nodes
 
@@ -925,6 +934,17 @@ def onehot_op(ctx, node, name, args):
         return [node, transpose_op]
     return node
 
+def fused_batchnorm_op7(ctx, node, name, args):
+    node.type = "BatchNormalization"
+    # tf inputs: x, scale, bias, mean, variance
+    # tf outputs: y, batch_mean, batch_var
+    # a: data_format, epsilon, is_training
+    # onnx inputs: X, scale, B, mean, variance, attributes: epsilon, momentum=0.9, spatial : 1
+    # output: mean, var, savedmean, savedvar,
+    nodes = conv_convert_inputs(ctx, node, with_kernel=False)
+    return nodes
+
+
 
 # pylint: enable=W0613,C0111,W0612
 
@@ -1048,6 +1068,7 @@ def onehot_op(ctx, node, name, args):
     "Sin": (direct_op, []),
     "Tan": (direct_op, []),
     "Multinomial": (multinomial_op, []),
+    "FusedBatchNorm": (fused_batchnorm_op7, []),
 }
 
 _OPSETS = [

tf2onnx/utils.py

@@ -76,7 +76,7 @@
     'dtype', 'output_shape', 'spatial', 'split', 'input_forget', 'keepdims', 'transA', 'auto_pad', 'border', 'low',
     'linear_before_reset', 'height_scale', 'output_padding', 'shape', 'kernel_shape', 'epsilon', 'size', 'starts',
     'direction', 'max', 'clip', 'across_channels', 'value', 'strides', 'extra_shape', 'scales', 'k', 'sample_size',
-    'blocksize'
+    'blocksize', 'epsilon', 'momentum'
 }