Merge pull request #315 from lucienwang1009/strided_slice

refine strided_slice

Author: pengwa

tests/test_backend.py

@@ -1242,6 +1242,28 @@ def test_strided_slice6(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
+    @unittest.skipIf(BACKEND in ["caffe2"], "multiple dims not supported")
+    def test_strided_slice7(self):
+        x_val = np.arange(5*6).astype("float32").reshape(5, 6)
+
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        x_ = tf.strided_slice(x, [0, 1], [3, 4], [1, 1], begin_mask=2)
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+
+        tf.reset_default_graph()
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        x_ = tf.strided_slice(x, [0, 1], [3, 4], [1, 1], end_mask=2)
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+
+        tf.reset_default_graph()
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        x_ = tf.strided_slice(x, [0, 1], [3, 4], [1, 1], shrink_axis_mask=2)
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+
+
     @unittest.skipIf(BACKEND in ["caffe2"], "fails with schema error")
     @unittest.skipIf(*support_op_conversion_since(7, "batchnorm"))
     def test_batchnorm(self):

tf2onnx/tfonnx.py

@@ -998,11 +998,14 @@ def stridedslice_op(ctx, node, name, args):
         if attr is not None and attr.i != 0:
             raise ValueError("StridedSlice: attribute " + attr_name + " not supported")
     input_shape = ctx.get_shape(node.input[0])
-    begin = node.inputs[1].get_tensor_value()
-    end = node.inputs[2].get_tensor_value()
-    strides = node.inputs[3].get_tensor_value()
+    begin = node.inputs[1].get_tensor_value(as_list=False)
+    end = node.inputs[2].get_tensor_value(as_list=False)
+    strides = node.inputs[3].get_tensor_value(as_list=False)
+    max_size = np.iinfo(begin.dtype).max
     end_mask = node.get_attr("end_mask")
     end_mask = end_mask.i if end_mask is not None else 0
+    begin_mask = node.get_attr("begin_mask")
+    begin_mask = begin_mask.i if begin_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 = []
@@ -1018,19 +1021,25 @@ def stridedslice_op(ctx, node, name, args):
 
         # an implicit condition is stride == 1 (checked in above)
         if begin_item < 0 and end_item == 0:
-            end_item = sys.maxsize
+            end_item = max_size
 
         mask = (shrink_axis_mask >> idx) & 1
         if mask != 0:
             new_begin.append(begin_item)
+            end_item = begin_item + 1 if begin_item != -1 else max_size
             new_end.append(end_item)
             needs_squeeze.append(idx)
             continue
 
-        new_begin.append(begin_item)
+        mask = (begin_mask >> idx) & 1
+        if mask != 0:
+            new_begin.append(0)
+        else:
+            new_begin.append(begin_item)
+
         mask = (end_mask >> idx) & 1
         if mask != 0:
-            new_end.append(sys.maxsize)
+            new_end.append(max_size)
         else:
             new_end.append(end_item)