Merge pull request #464 from zhijxu-MS/push_branch

support slice-10

Author: nbcsm

tests/test_backend.py

@@ -945,6 +945,42 @@ def test_slice(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
+    @check_opset_min_version(10, "Slice in opset 10 can accept dymaic 'start' and 'ends'")
+    def test_slice_with_non_const(self):
+        x_val = np.array([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=np.float32)
+        t1 = np.array([0, 1], dtype=np.int32)
+        t2 = np.array([2, 2], dtype=np.int32)
+        x0 = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        t1_ = tf.placeholder(tf.int32, t1.shape, name=_TFINPUT1)
+        t2_ = tf.placeholder(tf.int32, t2.shape, name=_TFINPUT2)
+        x_ = tf.slice(x0, t1_, t2_)
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val, _INPUT1: t1, _INPUT2: t2})
+
+    @check_opset_min_version(10, "Slice in opset 10 can accept dymaic 'start' and 'ends'")
+    def test_slice_with_size_is_negative_one(self):
+        x_val = np.array([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=np.float32)
+        t1 = np.array([0, 1], dtype=np.int32)
+        # input "size" contains -1
+        t2 = np.array([2, -1], dtype=np.int32)
+        x0 = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        t1_ = tf.placeholder(tf.int32, t1.shape, name=_TFINPUT1)
+        t2_ = tf.placeholder(tf.int32, t2.shape, name=_TFINPUT2)
+        x_ = tf.slice(x0, t1_, t2_)
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val, _INPUT1: t1, _INPUT2: t2})
+
+    @skip_caffe2_backend()
+    def test_slice1(self):
+        # FIXME: only 1 dimension supported by caffe2 and msrt
+        x_val = np.array([[[1, 1, 1], [2, 2, 2]], [[3, 3, 3], [4, 4, 4]], [[5, 5, 5], [6, 6, 6]]], dtype=np.float32)
+        t1 = tf.constant([1, 0, 0], dtype=tf.int32)
+        t2 = tf.constant([1, 1, 3], dtype=tf.int32)
+        x0 = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        x_ = tf.slice(x0, t1, t2)
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+
     def test_split(self):
         x_val = np.linspace(1.0, 5 * 30.0, 5 * 30).astype(np.float32).reshape(5, 30)
         x0 = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
@@ -1103,17 +1139,6 @@ def test_reducemean(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-    @unittest.skip("")
-    def test_slice1(self):
-        # FIXME: only 1 dimension supported by caffe2 and msrt
-        x_val = np.array([[[1, 1, 1], [2, 2, 2]], [[3, 3, 3], [4, 4, 4]], [[5, 5, 5], [6, 6, 6]]], dtype=np.float32)
-        t1 = tf.constant([1, 0, 0], dtype=tf.int32)
-        t2 = tf.constant([1, 1, 3], dtype=tf.int32)
-        x0 = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
-        x_ = tf.slice(x0, t1, t2)
-        _ = tf.identity(x_, name=_TFOUTPUT)
-        self._run_test_case([_OUTPUT], {_INPUT: x_val})
-
     @skip_caffe2_backend()
     @check_onnxruntime_incompatibility("Pow")
     def test_pow_scalar(self):
@@ -2078,5 +2103,6 @@ def test_space_to_batchnd(self):
         _ = tf.space_to_batch_nd(input_x, block_size, pad, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: input_val})
 
+
 if __name__ == '__main__':
     unittest_main()

tf2onnx/__init__.py

@@ -6,8 +6,8 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
-__all__ = ["utils", "graph_matcher", "graph", "loader", "tfonnx", "shape_inference", "schemas"]
+__all__ = ["utils", "graph_matcher", "graph", "graph_builder", "loader", "tfonnx", "shape_inference", "schemas"]
 
 from .version import version as __version__
 from . import verbose_logging as logging
-from tf2onnx import tfonnx, utils, graph, graph_matcher, shape_inference, schemas  # pylint: disable=wrong-import-order
+from tf2onnx import tfonnx, utils, graph, graph_builder, graph_matcher, shape_inference, schemas  # pylint: disable=wrong-import-order

tf2onnx/graph.py

@@ -20,6 +20,7 @@
 from tf2onnx.utils import port_name, find_opset
 from tf2onnx import optimizer
 from tf2onnx.schemas import get_schema, infer_onnx_shape_dtype
+from tf2onnx import constants
 
 logger = logging.getLogger(__name__)
 
@@ -419,7 +420,8 @@ def make_const(self, name, np_val, skip_conversion=False, raw=True):
         return node
 
     def make_node(self, op_type, inputs, attr=None, output_count=1, outputs=None, skip_conversion=True,
-                  op_name_scope=None, name=None, shapes=None, dtypes=None, domain=None, infer_shape_dtype=True):
+                  op_name_scope=None, name=None, shapes=None, dtypes=None, domain=constants.ONNX_DOMAIN,
+                  infer_shape_dtype=True):
         """Make a new onnx node in the graph"""
         if attr is None:
             attr = {}

tf2onnx/graph_builder.py

@@ -0,0 +1,104 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""
+tf2onnx.graph_helper - class to help building graph, such as helping to make complex node
+"""
+
+import numpy as np
+from tf2onnx import utils, logging
+
+
+# pylint: disable=missing-docstring
+
+
+logger = logging.getLogger(__name__)
+
+
+class GraphBuilder(object):
+    """help to build graph"""
+    def __init__(self, graph):
+        self._g = graph
+
+    @property
+    def graph(self):
+        return self._g
+
+    def make_slice(self, kwargs, name=None, shapes=None, dtypes=None):
+        """
+        slice changes its schema at opset 10: it treats some attributes as dynamic input
+        so this function has to process inputs according to graph's opset version
+        to get "inputs" and "attr" to feed "make_node"
+        kwargs: key could be ["data", "starts", "ends", "axes", "steps", "outputs"].
+        """
+        outputs = kwargs.pop("outputs", None)
+
+        if self.graph.opset < 10:
+            # "data" is string
+            # "starts", "ends" and "axes" are attributes, and "axes" is optional.
+            inputs = [kwargs.pop("data")]
+            starts = self.convert_to_attribute(kwargs.pop("starts"))
+            ends = self.convert_to_attribute(kwargs.pop("ends"))
+            axes = self.convert_to_attribute(kwargs.pop("axes", None), is_optional=True)
+            attr = {"starts": starts, "ends": ends, "axes": axes}
+        else:
+            # slice-10 has 3 required inputs "data", "starts", "ends"l
+            # and 2 optional inputs "axes", "steps"
+            # input sequence should be "data", "starts", "ends", "axes", "steps"
+            attr = {}
+            data = self.convert_to_input(kwargs.pop("data"))
+            starts = self.convert_to_input(kwargs.pop("starts"))
+            ends = self.convert_to_input(kwargs.pop("ends"))
+            axes = self.convert_to_input(kwargs.pop("axes", None), is_optional=True)
+            steps = self.convert_to_input(kwargs.pop("steps", None), is_optional=True)
+            inputs = [data, starts, ends, axes, steps]
+
+        # pro-process inputs and attr
+        if kwargs:
+            logger.warning("kwargs contains un-used key")
+
+        new_attr = {}
+        for key, val in attr.items():
+            if val is not None:
+                new_attr[key] = val
+        attr = new_attr
+
+        for ind, val in enumerate(inputs):
+            if val is None:
+                inputs[ind] = ""  # empty string means no connection in ONNX
+        # remove tailing ""
+        while inputs[-1] == "":
+            inputs = inputs[:-1]
+
+        return self.graph.make_node(op_type="Slice", inputs=inputs, attr=attr, name=name,
+                                    outputs=outputs, shapes=shapes, dtypes=dtypes).output[0]
+
+    def convert_to_input(self, tensor, is_optional=False):
+        """in ONNX, input shold come from node, so it must be a string"""
+        if is_optional and tensor is None:
+            return None
+
+        utils.make_sure(tensor is not None, "input is required so it couldn't be None")
+
+        res = tensor
+        if isinstance(tensor, list):
+            res = self.graph.make_const(utils.make_name("const_slice"), np.array(tensor)).output[0]
+
+        utils.make_sure(isinstance(res, str), "input is a dynamic input, so a str is needed")
+
+        return res
+
+    def convert_to_attribute(self, tensor, is_optional=False):
+        if is_optional and tensor is None:
+            return None
+
+        utils.make_sure(tensor is not None, "input is required so it couldn't be None")
+
+        res = tensor
+        if isinstance(tensor, str):
+            const_node = self.graph.get_node_by_output(tensor)
+            res = const_node.get_tensor_value(as_list=True)
+
+        utils.make_sure(isinstance(res, list), "input is an attr, so a list is needed")
+
+        return res

tf2onnx/onnx_opset/nn.py

@@ -15,6 +15,7 @@
 from onnx import onnx_pb
 from onnx.onnx_pb import TensorProto
 from tf2onnx import constants, utils
+from tf2onnx.graph_builder import GraphBuilder
 from tf2onnx.handler import tf_op
 from tf2onnx.onnx_opset import common, controlflow, tensor
 
@@ -463,6 +464,14 @@ def version_7(cls, ctx, node, **kwargs):
 
     @classmethod
     def version_9(cls, ctx, node, **kwargs):
+        cls._convert_since_9(ctx, node, **kwargs)
+
+    @classmethod
+    def version_10(cls, ctx, node, **kwargs):
+        cls._convert_since_9(ctx, node, **kwargs)
+
+    @classmethod
+    def _convert_since_9(cls, ctx, node, **kwargs):
         # float32 out = ResizeBilinear/ResizeNearestNeighbor(T images, int size)
         # https://www.tensorflow.org/api_docs/python/tf/image/resize_nearest_neighbor
         # wants the input to be NHWC - adjust target_shape to this.
@@ -481,8 +490,10 @@ def version_9(cls, ctx, node, **kwargs):
             scales = ctx.make_const(utils.make_name("scales"), scale_val, raw=False)
         else:
             ori_shape = ctx.make_node("Shape", [node.input[0]])
-            ori_shape_hw = ctx.make_node("Slice", ori_shape.output, {"axes": [0], "starts": [1], "ends": [3]})
-            ori_shape_hw_float = ctx.make_node("Cast", ori_shape_hw.output, attr={"to": onnx_pb.TensorProto.FLOAT})
+            attr = {"axes": [0], "starts": [1], "ends": [3]}
+            inputs_map = {"data": ori_shape.output[0], **attr}
+            ori_shape_hw = GraphBuilder(ctx).make_slice(inputs_map)
+            ori_shape_hw_float = ctx.make_node("Cast", [ori_shape_hw], attr={"to": onnx_pb.TensorProto.FLOAT})
 
             target_hw = node.inputs[1]
             target_hw_float = ctx.make_node("Cast", target_hw.output, attr={"to": onnx_pb.TensorProto.FLOAT})
@@ -538,12 +549,13 @@ def version_7(cls, ctx, node, **kwargs):
         new_line = g.make_node(op_type="Concat", inputs=[const_zero_bool.output[0], "line"],
                                attr={"axis": counter_axis},
                                dtypes=[onnx_pb.TensorProto.BOOL])
-        slice_node = g.make_node(op_type="Slice", inputs=[new_line.output[0]],
-                                 attr={"axes": [counter_axis], "starts": [0], "ends": [-1]})
+        attr = {"axes": [counter_axis], "starts": [0], "ends": [-1]}
+        inputs_map = {"data": new_line.output[0], **attr}
+        slice_node = GraphBuilder(g).make_slice(inputs_map)
 
         g.make_node("Identity", ["cond"], outputs=["cond_out"])
         g.make_node("Identity", ["line"], outputs=["res"])
-        g.make_node("Identity", [slice_node.output[0]], outputs=["line_out"])
+        g.make_node("Identity", [slice_node], outputs=["line_out"])
 
         g.add_graph_input("trip", onnx_pb.TensorProto.INT64, [])
         g.add_graph_input("cond", onnx_pb.TensorProto.BOOL, [])