support for more complex strided_slice options

Author: guschmue

tests/test_backend.py

@@ -69,9 +69,6 @@ def get_conv_getdata(kind=1):
         data = [
             ('SAME', [32, 35, 35, 288], [1, 3, 3, 1], [1, 2, 2, 1]),
             ('SAME', [32, 35, 35, 288], [1, 2, 2, 1], [1, 2, 2, 1]),
-            ('SAME', [32, 35, 35, 288], [1, 2, 2, 1], [1, 1, 1, 1]),
-            ('SAME', [32, 35, 35, 288], [1, 5, 5, 1], [1, 1, 1, 1]),
-            ('SAME', [32, 35, 35, 288], [1, 1, 1, 1], [1, 2, 2, 1]),
             ('SAME', [32, 35, 35, 288], [1, 1, 1, 1], [1, 1, 1, 1]),
             ('SAME', [32, 35, 35, 288], [1, 5, 2, 1], [1, 2, 2, 1]),
             ('SAME', [32, 35, 35, 288], [1, 2, 5, 1], [1, 2, 2, 1]),
@@ -83,7 +80,6 @@ def get_conv_getdata(kind=1):
             ('SAME', [1, 28, 28, 3], [1, 5, 5, 1], [1, 2, 2, 1]),
             ('SAME', [1, 28, 28, 3], [1, 5, 5, 1], [1, 1, 1, 1]),
             ('SAME', [1, 28, 28, 3], [1, 5, 2, 1], [1, 2, 2, 1]),
-            ('SAME', [1, 28, 28, 3], [1, 2, 5, 1], [1, 1, 1, 1]),
             ('SAME', [32, 8, 8, 2048], [1, 3, 3, 1], [1, 2, 2, 1]),
             ('SAME', [32, 8, 8, 2048], [1, 3, 3, 1], [1, 1, 1, 1]),
             ('VALID', [32, 35, 35, 288], [1, 3, 3, 1], [1, 1, 1, 1]),
@@ -975,7 +971,6 @@ def test_strided_slice1(self):
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
 
-    @unittest.skipIf(BACKEND in ["caffe2", "onnxmsrt"], "multiple dims not supported")
     def test_strided_slice2(self):
         x_val = np.arange(3*2*3).astype("float32").reshape(3, 2, 3)
         x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
@@ -984,7 +979,6 @@ 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)
@@ -993,7 +987,6 @@ def test_strided_slice3(self):
         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)
@@ -1002,7 +995,7 @@ def test_strided_slice4(self):
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
 
-    @unittest.skip
+    @unittest.skipIf(BACKEND in ["caffe2", "onnxmsrt"], "multiple dims not supported")
     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)
@@ -1011,6 +1004,17 @@ def test_strided_slice5(self):
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
 
+    @unittest.skipIf(BACKEND in ["caffe2", "onnxmsrt"], "multiple dims not supported")
+    def test_strided_slice6(self):
+        # example from here:
+        # https://www.tensorflow.org/versions/r1.0/api_docs/cc/class/tensorflow/ops/strided-slice
+        x_val = np.arange(5*6).astype("float32").reshape(5, 6)
+        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.skipIf(BACKEND in ["caffe2", "onnxmsrt"], "fails with schema error")
     def test_batchnorm(self):
         x_shape = [1, 28, 28, 2]
@@ -1075,7 +1079,7 @@ def test_fill(self):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument('--backend', default='caffe2',
+    parser.add_argument('--backend', default=BACKEND,
                         choices=["caffe2", "onnxmsrt", "onnxmsrtnext", "onnx-tensorflow"],
                         help="backend to test against")
     parser.add_argument('--opset', default=OPSET,

tf2onnx/graph.py

@@ -31,7 +31,7 @@ def __init__(self, node, graph):
         for a in node.attribute:
             self._attr[a.name] = a
         # try to find a dtype for this node
-        dtype = graph.dtypes.get(node.name)
+        dtype = graph._dtypes.get(node.name)
         if not dtype:
             dtype = self._attr.get("dtype")
             if dtype:
@@ -105,9 +105,10 @@ def __str__(self):
     def __repr__(self):
         return "<onnx op type='%s' name=%s>" % (self.type, self._op.name)
 
-    def get_attr(self, name):
+    def get_attr(self, name, default=None):
         """Get attribute map."""
-        return self.attr.get(name)
+        attr = self.attr.get(name, default)
+        return attr
 
     def set_attr(self, name, value):
         self.attr[name] = helper.make_attribute(name, value)
@@ -236,7 +237,7 @@ def __init__(self, nodes, output_shapes=None, dtypes=None, target=None, opset=No
         self.shapes = {}
         self.model_inputs = []
         self._target = set(target)
-        self.dtypes = dtypes
+        self._dtypes = dtypes
         self._output_shapes = output_shapes
         ops = [Node(node, self) for node in nodes]
         self.set_nodes(ops)
@@ -290,6 +291,10 @@ def add_initializer(self, tensor):
         """Add tensor to initializers."""
         self._initializers[tensor.name] = tensor
 
+    def get_dtype(self, name):
+        """Get dtype for node."""
+        return self._dtypes.get(name)
+
     def get_shape(self, name):
         """Get shape for node."""
         assert isinstance(name, str)

tf2onnx/tfonnx.py

@@ -733,24 +733,72 @@ def expanddims_op7(ctx, node, name, args):
 
 
 def stridedslice_op(ctx, node, name, args):
-    # only the cases strides=1 can be mapped to onnx
-    not_supported_attr = ["begin_mask", "ellipsis_mask", "end_mask", "new_axis_mask", "shrink_axis_mask"]
+    # for now we implement common cases. Things like strides!=1 are not mappable to onnx.
+    not_supported_attr = ["ellipsis_mask", "new_axis_mask"]
     for attr_name in not_supported_attr:
         attr = node.get_attr(attr_name)
         if attr is not None and attr.i != 0:
-            raise ValueError("StridedSlice: attribute " + attr_name + " must be 0")
+            raise ValueError("StridedSlice: attribute " + attr_name + " not supported")
+
     begin = node.inputs[1].get_tensor_value()
     end = node.inputs[2].get_tensor_value()
-    strides = node.inputs[3].get_tensor_value()[0]
-    if strides != 1:
-        raise ValueError("StridedSlice: only strides=1 is supported")
-    node.set_attr("starts", list(begin))
-    node.set_attr("ends", list(end))
+    strides = node.inputs[3].get_tensor_value()
+    end_mask = node.get_attr("end_mask")
+    end_mask = end_mask.i if end_mask is not None else 0
+    shrink_axis_mask = node.get_attr("shrink_axis_mask")
+    shrink_axis_mask = shrink_axis_mask.i if shrink_axis_mask is not None else 0
+    new_begin = []
+    new_end = []
+    axes = []
+    # onnx slice op can't remove a axis, track axis and add a squeeze op if needed
+    needs_squeeze = []
+    for idx in range(len(begin)):
+        if strides[idx] != 1:
+            raise ValueError("StridedSlice: only strides=1 is supported")
+        axes.append(idx)
+        mask = (shrink_axis_mask >> idx) & 1
+        if mask != 0:
+            new_begin.append(begin[idx])
+            new_end.append(end[idx])
+            needs_squeeze.append(idx)
+            continue
+
+        new_begin.append(begin[idx])
+        mask = (end_mask >> idx) & 1
+        if mask != 0:
+            new_end.append(sys.maxsize)
+        else:
+            new_end.append(end[idx])
+
+    node.set_attr("starts", new_begin)
+    node.set_attr("ends", new_end)
+    node.set_attr("axes", axes)
     node.type = "Slice"
     ctx.remove_input(node, node.input[3])
     ctx.remove_input(node, node.input[2])
     ctx.remove_input(node, node.input[1])
-    return node
+    nodes = [node]
+    if needs_squeeze:
+        name = utils.make_name(node.name)
+        squeeze_op = ctx.insert_new_node_on_output("Squeeze", node.output[0], name)
+        squeeze_op.set_attr("axes", needs_squeeze)
+        nodes.append(squeeze_op)
+        ctx.copy_shape(node.output[0], squeeze_op.output[0])
+
+    # onnx slice as of opset 7 does only take float tensors ... cast if needed
+    input_dtype = ctx.get_dtype(node.input[0])
+    if input_dtype in [onnx_pb.TensorProto.INT32, onnx_pb.TensorProto.INT64]:
+        cast_op = ctx.insert_new_node_on_input(node, "Cast", node.input[0])
+        cast_op.set_attr("to", onnx_pb.TensorProto.FLOAT)
+        ctx.copy_shape(node.input[0], cast_op.output[0])
+        nodes.insert(0, cast_op)
+        name = utils.make_name(node.name)
+        cast_op = ctx.insert_new_node_on_output("Cast", nodes[-1].output[0], name)
+        cast_op.set_attr("to", input_dtype)
+        ctx.copy_shape(node.input[0], cast_op.output[0])
+        nodes.append(cast_op)
+
+    return nodes
 
 
 def pow_op(ctx, node, name, args):
@@ -1267,7 +1315,7 @@ def tensorflow_onnx_mapping(g, continue_on_error, custom_op_handlers):
 def tf_optimize(sess, inputs, outputs, graph_def):
     """Optimize tensorflow graph for inference."""
     transforms = [
-        "fold_constants(ignore_errors=true)",
+        #"fold_constants(ignore_errors=true)",
         "fold_batch_norms",
         "fold_old_batch_norms",
     ]