ONNX export: Slice op - Handle None value for ends (#14942)

* ONNX export: Slice op - Handle None value for ends
apache · Feb 20, 2020 · 9dcf71d · 9dcf71d
1 parent ab48a43
commit 9dcf71d
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 5 deletions.
diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
@@ -1495,17 +1495,20 @@ def convert_slice_axis(node, **kwargs):
 
     axes = int(attrs.get("axis"))
     starts = int(attrs.get("begin"))
-    ends = int(attrs.get("end", None))
-    if not ends:
-        raise ValueError("Slice: ONNX doesnt't support 'None' in 'end' attribute")
+    ends = attrs.get("end", None)
+    if not ends or ends == 'None':
+        # ONNX doesn't support None for ends. Since ends=None depicts
+        # length of dimension, passing dimension in this case.
+        in_shape = kwargs['in_shape'][0]
+        ends = in_shape[axes]
 
     node = onnx.helper.make_node(
         "Slice",
         input_nodes,
         [name],
         axes=[axes],
         starts=[starts],
-        ends=[ends],
+        ends=[int(ends)],
         name=name,
     )
     return [node]

diff --git a/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py b/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py
@@ -499,15 +499,19 @@ def split(attrs, inputs, proto_obj):
 
 def _slice(attrs, inputs, proto_obj):
     """Returns a slice of the input tensor along multiple axes."""
+    input_tensor_data = proto_obj.model_metadata.get('input_tensor_data')[0]
+    input_shape = input_tensor_data[1]
     new_attrs = translation_utils._fix_attribute_names(attrs,
                                                        {'axes' : 'axis',
                                                         'ends' : 'end',
                                                         'starts' : 'begin'})
     # onnx slice provides slicing on multiple axis. Adding multiple slice_axis operator
     # for multiple axes from mxnet
     begin = new_attrs.get('begin')
-    end = new_attrs.get('end')
+    end = list(new_attrs.get('end'))
     axes = new_attrs.get('axis', tuple(range(len(begin))))
+    for i, axis in enumerate(axes):
+        end[i] = None if end[i] >= input_shape[axis] else end[i]
     slice_op = symbol.slice_axis(inputs[0], axis=axes[0], begin=begin[0], end=end[0])
     if len(axes) > 1:
         for i, axis in enumerate(axes):

diff --git a/tests/python-pytest/onnx/test_cases.py b/tests/python-pytest/onnx/test_cases.py
@@ -41,6 +41,7 @@
              'test_globalaveragepool',
              'test_slice_cpu',
              'test_slice_neg',
+             'test_slice_end',
              'test_reciprocal',
              'test_sqrt',
              'test_pow',