[Relay] Add ResizeNearestNeighbor and CropAndResize in tf converter

Author: yongwww

python/tvm/relay/frontend/tensorflow.py

@@ -484,6 +484,60 @@ def _impl(inputs, attr, params):
         return inputs[0]
     return _impl
 
+def _crop_and_resize():
+    def _impl(inputs, attr, params):
+        # input image is a 4-D tensor of shape [batch, image_height, image_width, depth]
+        # boxes is a 2-D tensor of shape [num_boxes, 4], 4 is for [y1, x1, y2, x2]
+        boxes = params.pop(inputs[1].name_hint).asnumpy()
+        data_shape = attr['_input_shapes'][inputs[0]]
+        data_dim = len(data_shape)
+        box_ind = params.pop(inputs[2].name_hint).asnumpy().tolist()
+        crop_size = params.pop(inputs[3].name_hint).asnumpy().tolist()
+        method = attr['method'].decode()
+
+        # 1) Crop
+        # y is mapped to the image coordinate at y * (image_height - 1)
+        # x is mapped to the image coordinate at x * (image_width - 1)
+        begin = [0] * data_dim
+        begin[1] = float(boxes[0][0]) * (data_shape[1] - 1)
+        begin[2] = int(round(boxes[0][1] * (data_shape[2] - 1)))
+        size = data_shape[:]
+        size[0] = 1
+        size[1] = int(round((data_shape[1] - 1) * boxes[0][2])) + 1
+        size[2] = int(round((data_shape[2] - 1) * boxes[0][3])) + 1
+        res_crop = _op.strided_slice(inputs[0], begin=begin, end=size)
+
+        # 2) Resize
+        attrs = {};
+        attrs['size'] = crop_size
+        attrs['layout'] = 'NHWC'
+        if method.lower() == 'nearest':
+            raise tvm.error.OpAttributeUnimplemented(
+                'Attribute method=nearest is not supported')
+        else:
+            attrs['align_corners'] = True
+            attrs['method'] = 'BILINEAR'
+        ret = _get_relay_op('resize')(res_crop, **attrs)
+
+        for idx in box_ind[1:]:
+            # 1) Crop
+            # y is mapped to the image coordinate at y * (image_height - 1)
+            begin = [0] * data_dim
+            begin[0] = idx
+            begin[1] = int(round(boxes[idx][0] * (data_shape[1] - 1)))
+            begin[2] = int(round(boxes[idx][1] * (data_shape[2] - 1)))
+            size = data_shape[:]
+            size[0] = idx + 1
+            size[1] = int(round((data_shape[1] - 1) * boxes[idx][2])) + 1
+            size[2] = int(round((data_shape[2] - 1) * boxes[idx][3])) + 1
+            res_crop = _op.strided_slice(inputs[0], begin=begin, end=size)
+
+            # 2) Resize
+            res_resize = _get_relay_op('resize')(res_crop, **attrs)
+            ret = _op.concatenate([ret, res_resize], axis=0)
+        return ret
+    return _impl
+
 def _cast():
     def _impl(inputs, attr, params):
         return inputs[0].astype(attr['DstT'].name)
@@ -514,6 +568,21 @@ def _impl(inputs, attr, params):
                        extras={'method': "BILINEAR"})(inputs, attr)
     return _impl
 
+def _resize_nearest_neighbor():
+    def _impl(inputs, attr, params):
+        size = attr['_output_shapes'][0][1:3]
+        if -1 in size:
+            size = _infer_value(inputs[1], params).asnumpy().reshape([-1]).tolist()
+        attr['size'] = size
+        inputs.pop(1)
+        # NHWC
+        attr['layout'] = 'NHWC'
+
+        return AttrCvt(op_name="resize",
+                       ignores=['Tdim'],
+                       extras={'method': "NEAREST_NEIGHBOR"})(inputs, attr)
+    return _impl
+
 def _check_numerics():
     def _impl(inputs, attr, params):
         # Making a copy node assuming no need to verify
@@ -593,7 +662,7 @@ def _impl(inputs, attr, params):
                 end[i] = data_shape[i] - begin[i]
             else:
                 end[i] += begin[i]
-        return _op.strided_slice(inputs[0], begin=begin, end=size)
+        return _op.strided_slice(inputs[0], begin=begin, end=end)
     return _impl
 
 
@@ -1243,6 +1312,7 @@ def _impl(inputs, attr, params):
     'Concat'                            : _concat(),
     'ConcatV2'                          : _concatV2(),
     'Conv2D'                            : _conv('conv'),
+    'CropAndResize'                     : _crop_and_resize(),
     'DecodeJpeg'                        : _decode_image(),
     'DepthwiseConv2dNative'             : _conv('depthwise'),
     'DepthToSpace'                      : _depth_to_space(),
@@ -1295,6 +1365,7 @@ def _impl(inputs, attr, params):
     'Reshape'                           : _reshape(),
     'ResizeBilinear'                    : _resize_bilinear(),
     'ResizeBicubic'                     : _resize_bilinear(),
+    'ResizeNearestNeighbor'             : _resize_nearest_neighbor(),
     'ReverseV2'                         : _reverse_v2(),
     'RightShift'                        : AttrCvt('right_shift'),
     'Round'                             : AttrCvt('round'),

python/tvm/relay/frontend/tensorflow_parser.py

@@ -18,7 +18,6 @@
 from __future__ import absolute_import as _abs
 from __future__ import print_function
 import os
-from tensorflow.core.framework import graph_pb2
 from tvm.contrib import util
 
 
@@ -35,12 +34,12 @@ class TFParser(object):
     --------
     .. code-block:: python
 
-        parser = TfParser(model_dir)
-        graph = parser.parse()
-        # graph is related graphdef of the model
+        parser = TFParser(model_dir)
+        graphdef = parser.parse()
     """
 
     def __init__(self, model_dir):
+        from tensorflow.core.framework import graph_pb2
         self._tmp_dir = util.tempdir()
         self._model_dir = model_dir
         self._graph = graph_pb2.GraphDef()
@@ -96,6 +95,7 @@ def _load_saved_model(self):
             from tensorflow.python.tools import freeze_graph
             from tensorflow.python.framework import ops
             from tensorflow.python.framework import graph_util
+            from tensorflow.core.framework import graph_pb2
         except ImportError:
             raise ImportError(
                 "InputConfiguration: Unable to import tensorflow which is "

tests/python/frontend/tensorflow/test_forward.py

@@ -969,8 +969,8 @@ def test_forward_multi_output():
                 tvm.testing.assert_allclose(tf_output[i], tvm_output[i], atol=1e-5, rtol=1e-5)
 
 #######################################################################
-# Resize Bilinear
-# ---------------
+# Resize Bilinear, Nearest_Neighbor
+# ---------------------------------
 
 def _test_resize_bilinear(in_shape, to_shape, align_corners):
     """ One iteration of resize bilinear """
@@ -1000,13 +1000,31 @@ def _test_resize_bilinear_from_tensor(in_shape, align_corners):
 
         compare_tf_with_tvm(data, 'Placeholder:0', 'ResizeBilinear:0')
 
-def test_forward_resize_bilinear():
-    """ Resize Bilinear """
+
+def _test_resize_nearest_neighbor(in_shape, to_shape):
+    """ One iteration of resize nearest neighbor """
+
+    data = np.random.uniform(size=in_shape).astype('float32')
+    shape_data = np.array(to_shape).astype('int32')
+
+    with tf.Graph().as_default():
+        in_data = array_ops.placeholder(shape=data.shape, dtype=data.dtype)
+        shape_data = constant_op.constant(
+            shape_data, shape=shape_data.shape, dtype=shape_data.dtype)
+        tf.image.resize_nearest_neighbor(in_data, shape_data, name='resize_nearest_neighbor')
+
+        compare_tf_with_tvm(data, 'Placeholder:0', 'resize_nearest_neighbor:0')
+
+
+def test_forward_resize():
+    """ Resize Bilinear, Nearest_Neighbor """
 
     _test_resize_bilinear((4, 16, 32, 32), [50, 50], False)
     _test_resize_bilinear((6, 32, 64, 64), [20, 20], True)
     _test_resize_bilinear_from_tensor((4, 16, 32, 32), False)
     _test_resize_bilinear_from_tensor((6, 32, 50, 50), True)
+    _test_resize_nearest_neighbor((6, 32, 64, 64), [20, 20])
+
 
 #######################################################################
 # BroadcastTo
@@ -1100,6 +1118,37 @@ def test_forward_crop():
     _test_crop((1, 224, 224, 3), 20, 20, 120, 120)
 
 
+#######################################################################
+# CropAndResize
+# -------------
+
+def _test_forward_crop_and_resize(img_shape, boxes, box_idx, crop_size, method='bilinear', dtype="float32"):
+    image = np.random.uniform(0, 10, size=img_shape).astype(dtype)
+    tf.reset_default_graph()
+    in_data = tf.placeholder(dtype, image.shape, name="in_data")
+    tf.image.crop_and_resize(in_data, boxes=boxes, box_ind=box_idx, crop_size=crop_size,
+                             method=method, name="crop_and_resize")
+    compare_tf_with_tvm([image], ['in_data:0'], 'crop_and_resize:0')
+
+def test_forward_crop_and_resize():
+    """ CropAndResize """
+    _test_forward_crop_and_resize([1, 11, 11, 3], [[0, 0, 1, 1]],
+                                  [0], [5, 5])
+    _test_forward_crop_and_resize([1, 11, 11, 3], [[0, 0, .9, .9]],
+                                  [0], [5, 5])
+    _test_forward_crop_and_resize([1, 11, 11, 3], [[.1, .2, 1, 1]],
+                                  [0], [5, 5])
+    _test_forward_crop_and_resize([1, 21, 21, 3], [[.2, .3, .7, .9]],
+                                  [0], [3, 4])
+    _test_forward_crop_and_resize([10, 11, 11, 3],
+                                  [[0, 0, 0.9, 0.9], [0.2, 0.2, 0.8, 0.8]],
+                                  [0, 1], [5, 5])
+    _test_forward_crop_and_resize([3, 11, 11, 3],
+                                  [[0, 0, 0.9, 0.9], [0.2, 0.2, 0.8, 0.8], [0, 0, 1, 1]],
+                                  [0, 1, 2], [3, 3])
+    _test_forward_crop_and_resize([1, 16, 16, 3], [[0.2, 0.4, 0.8, 0.8]], [0], [3, 3])
+
+
 #######################################################################
 # LSTM
 # ----
@@ -1989,10 +2038,11 @@ def test_placeholder():
     test_forward_depthtospace()
     test_forward_squeeze()
     test_forward_pack()
-    test_forward_resize_bilinear()
     test_forward_broadcast_to()
     test_forward_fill()
     test_forward_crop()
+    test_forward_resize()
+    test_forward_crop_and_resize()
     test_forward_pad()
     test_forward_unpack()
     test_forward_gather()