Dynamic ONNX importer: Upsampling and Pad (#2)

fix lint

fix Call reference

fix a type issue with expand

fix a bad test refactor

respond to review comments, fix batch matmul tests

Author: Lily Orth-Smith

include/tvm/relay/transform.h

@@ -213,7 +213,7 @@ TVM_DLL Pass FastMath();
  *
  * Searches the graph for dynamic ops. If the dynamic inputs to those ops are constants, it replaces
  * them with static ops and re-performs type inference and constant folding. The pass repeats
- * istself until the graph stops changing or we run too many iterations.
+ * itself until the graph stops changing or we run too many iterations.
  *
  * \return The pass.
  */

python/tvm/relay/frontend/onnx.py

@@ -28,17 +28,9 @@
 from .. import op as _op
 from .. import vision as _vision
 
-from ..function import Function
-from ..expr import Call, Let
-from ..expr import If, Tuple, TupleGetItem
-from ..expr import RefCreate, RefRead, RefWrite
-from ..expr_functor import ExprFunctor
-from ..adt import Match, Clause
-from ..op.tensor import minimum as _minimum, maximum as _maximum
-
 from .common import AttrCvt, Renamer
 from .common import get_relay_op, new_var, infer_shape, infer_channels
-from .common import infer_type, get_name, infer_value, infer_value_simulated
+from .common import infer_type, get_name, infer_value_simulated
 
 __all__ = ['from_onnx']
 
@@ -642,26 +634,22 @@ def _impl_v2(cls, inputs, attr, params):
 
     @classmethod
     def _impl_v11(cls, inputs, attr, params):
-        pad_width = []
-        pads = infer_value_simulated(inputs[1], params).asnumpy()
+        pads = inputs[1]
         if len(inputs) == 3:
-            value = infer_value_simulated(inputs[2], params).asnumpy().item()
+            value = _op.take(inputs[2], _op.const(0))
         else:
             value = 0
-        attr["pad_value"] = value
-        dims = int(len(pads) / 2)
-        for i in range(dims):
-            pad_width.append((pads[i], pads[i + dims]))
-        attr['pad_width'] = pad_width
+
+        pads_shape = infer_shape(pads)
+        dims = int(pads_shape[0] / 2)
+        pad_width_expr = _op.transpose(_op.reshape(pads, (2, dims)))
         pad_mode = attr.get('mode', b'constant').decode('utf-8')
-        if pad_mode in ['constant', 'edge', 'reflect']:
-            attr['pad_mode'] = pad_mode
-            attr.pop('mode', None)
-        else:
+
+        if not pad_mode in ['constant', 'edge', 'reflect']:
             raise tvm.error.OpAttributeInvalid('Value ' + pad_mode +
                                                ' in attribute "mode" is invalid for operator Pad.')
 
-        return AttrCvt('pad')(inputs[:1], attr, params)
+        return _op.nn.pad(inputs[0], pad_width_expr, value, pad_mode=pad_mode)
 
 
 class ParametricSoftPlus(OnnxOpConverter):
@@ -869,17 +857,24 @@ class Upsample(OnnxOpConverter):
     @classmethod
     def _impl_v9(cls, inputs, attr, params):
         scales = attr.get('scales')
+
+        input_shape = infer_shape(inputs[0])
+        dims = len(input_shape)
+
         if not scales:
             #Here we are going to higher OPSET version.
-            assert len(inputs) == 2, "Upsample op take 2 inputs, {} given".format(len(inputs))
+            assert len(inputs) == 2, "Upsample op takes 2 inputs, {} given".format(len(inputs))
+
             if get_name(inputs[1]) in params:
                 scales = params[inputs[1].name_hint].asnumpy()
-            else:
+            elif dims == 5:
                 scales = infer_value_simulated(inputs[1], params).asnumpy()
-            inputs = inputs[:1]
-        assert scales[0] == 1.0 and scales[1] == 1.0
-        input_shape = infer_shape(inputs[0])
-        dims = len(input_shape)
+            else:
+                scales = inputs[1]
+
+        if not isinstance(scales, _expr.Call):
+            assert scales[0] == 1.0 and scales[1] == 1.0
+
         mode = attr.get('mode')
         if mode == b'nearest':
             method = "nearest_neighbor"
@@ -888,21 +883,41 @@ def _impl_v9(cls, inputs, attr, params):
         else:
             raise tvm.error.OpAttributeInvalid(
                 'Value {} in attribute "mode" of operator Upsample is not valid.'.format(mode))
-        attr = {'scale_h': scales[-2], 'scale_w': scales[-1], 'method': method}
+
+        if method == 'nearest_neighbor':
+            align_corners = False
+        else:
+            align_corners = True
+        # in 3d case, we use the purely static op
         if dims == 5:
-            assert len(scales) == 5
-            attr['scale_d'] = scales[-3]
-            attr['layout'] = 'NCDHW'
-            op_name = 'upsampling3d'
+            scale_h = scales[-2]
+            scale_w = scales[-1]
+            scale_d = scales[-3]
+            layout = 'NCDHW'
+            out = _op.nn.upsampling3d(inputs[0],
+                                      scale_d,
+                                      scale_h,
+                                      scale_w,
+                                      layout=layout,
+                                      method=method)
+        # in 2d case, use dynamic op
         else:
-            assert len(scales) == 4
-            attr['layout'] = 'NCHW'
-            if method == 'nearest_neighbor':
-                attr['align_corners'] = False
+            if isinstance(scales, _expr.Call):
+                scale_h = _op.take(scales, _op.const(3))
+                scale_w = _op.take(scales, _op.const(4))
             else:
-                attr['align_corners'] = True
-            op_name = 'upsampling'
-        return AttrCvt(op_name)(inputs, attr)
+                assert len(scales) == 4
+                scale_h = scales[-2]
+                scale_w = scales[-1]
+            layout = 'NCHW'
+
+            out = _op.nn.upsampling(inputs[0],
+                                    scale_h,
+                                    scale_w,
+                                    layout=layout,
+                                    method=method,
+                                    align_corners=align_corners)
+        return out
 
 
 class Shape(OnnxOpConverter):
@@ -1422,7 +1437,8 @@ class Expand(OnnxOpConverter):
     """
     @classmethod
     def _impl_v8(cls, inputs, attr, params):
-        in_shape = _op.shape_of(inputs[0])
+        dtype = infer_type(inputs[1]).checked_type.dtype
+        in_shape = _op.shape_of(inputs[0], dtype=dtype)
         shape = inputs[1]
 
         # Currently 'op.broadcast_to' expect the rank of the given 'shape'
@@ -1441,14 +1457,11 @@ def expand_shape(in_shape, shape):
             in_dims = infer_shape(in_shape)[0]
             new_dims = infer_shape(shape)[0]
             if in_dims < new_dims:
-                in_shape = _op.concatenate([_expr.const([
-                    1,
-                ] * (new_dims - in_dims)), in_shape],
-                                           axis=0)
+                in_shape = _op.concatenate([_expr.const([1, ] * (new_dims - in_dims), dtype=dtype),
+                                            in_shape], axis=0)
             elif new_dims > in_dims:
-                shape = _op.concatenate([_expr.const([
-                    1,
-                ] * (in_dims - new_dims)), shape], axis=0)
+                shape = _op.concatenate([_expr.const([1, ] * (in_dims - new_dims), dtype=dtype),
+                                         shape], axis=0)
             new_shape = _op.maximum(in_shape, shape)
             return new_shape
 
@@ -2058,6 +2071,13 @@ def from_onnx(self, graph, opset, freeze_params=False):
 
         opset : opset version
 
+        freeze_params: bool
+            If this parameter is true, the importer will take any provided
+            onnx input values (weights, shapes, etc) and embed them into the relay model
+            as Constants instead of variables. This allows more aggressive optimizations
+            at compile time and helps in making models static if certain inputs represent
+            attributes relay would traditionally consider compile-time constants.
+
         Returns
         -------
         mod : tvm.IRModule
@@ -2156,12 +2176,12 @@ def from_onnx(self, graph, opset, freeze_params=False):
         ## Maintain the order of inputs and parametersfrom the ONNX graph, but only include
         ## those parameters that are needed to execute the relay graph
         free_vars = analysis.free_vars(outputs)
-        nodes = {v:k for k,v in self._nodes.items()}
+        nodes = {v: k for k, v in self._nodes.items()}
         free_vars = [nodes[var] for var in free_vars]
         for i_name in self._params:
             if i_name in free_vars and i_name not in self._inputs:
                 self._inputs[i_name] = self._nodes[i_name]
-        func = _function.Function([v for k,v in self._inputs.items()], outputs)
+        func = _function.Function([v for k, v in self._inputs.items()], outputs)
         if freeze_params:
             func, params = self.freeze(func, self._params)
             return IRModule.from_expr(func), params
@@ -2282,6 +2302,13 @@ def from_onnx(model, shape=None, dtype="float32", opset=None, freeze_params=Fals
         Override to autodetected opset.
         This can be helpful for some testing.
 
+    freeze_params: bool
+        If this parameter is true, the importer will take any provided
+        onnx input values (weights, shapes, etc) and embed them into the relay model
+        as Constants instead of variables. This allows more aggressive optimizations
+        at compile time and helps in making models static if certain inputs represent
+        attributes relay would traditionally consider compile-time constants.
+
     Returns
     -------
     mod : tvm.IRModule

python/tvm/topi/x86/batch_matmul.py

@@ -50,6 +50,10 @@ def batch_matmul(cfg, x, y, out_shape=None):
     assert XK == YK, "shapes of x and y is inconsistant"
     B = XB
     K = XK
+    if out_shape is not None:
+        assert out_shape[0] == B, "got invalid output shape"
+        assert out_shape[1] == M, "got invalid output shape"
+        assert out_shape[2] == N, "got invalid output shape"
     if cfg.is_fallback:
         _default_batch_matmul_config(cfg, M, N, K)
 

tests/python/frontend/onnx/test_forward.py

@@ -48,7 +48,7 @@ def get_tvm_output_with_vm(graph_def, input_data, target, ctx, opset=None, freez
     """ Generic function to execute and get tvm output with vm executor"""
     if not isinstance(input_data, list):
         input_data = [input_data]
-    input_names, shape_dict = get_input_data_shape_dict(graph_def, input_data)
+    _, shape_dict = get_input_data_shape_dict(graph_def, input_data)
 
     mod, params = relay.frontend.from_onnx(graph_def, shape_dict, opset=opset, freeze_params=freeze_params)
 
@@ -167,15 +167,26 @@ def test_reshape():
 # @tvm.testing.uses_gpu
 def test_expand():
 
-    def _test_expand(name, data, shape, ref_data):
+    def _test_expand(name, data, shape, ref_data, dtype="int32"):
         shape_array = np.array(shape)
-        shape_node = onnx.helper.make_node('Constant',
-                                    inputs=[],
-                                    outputs=['shape'],
-                                    value=onnx.helper.make_tensor(name = 'const_tensor',
-                                                                  data_type = onnx.TensorProto.INT32,
-                                                                  dims = shape_array.shape,
-                                                                  vals = shape_array.flatten().astype('int32')))
+        if dtype == "int32":
+            shape_node = onnx.helper.make_node('Constant',
+                                        inputs=[],
+                                        outputs=['shape'],
+                                        value=onnx.helper.make_tensor(name = 'const_tensor',
+                                                                      data_type = onnx.TensorProto.INT32,
+                                                                      dims = shape_array.shape,
+                                                                      vals = shape_array.flatten().astype('int32')))
+        elif dtype == "int64":
+            shape_node = onnx.helper.make_node('Constant',
+                                        inputs=[],
+                                        outputs=['shape'],
+                                        value=onnx.helper.make_tensor(name = 'const_tensor',
+                                                                      data_type = onnx.TensorProto.INT64,
+                                                                      dims = shape_array.shape,
+                                                                      vals = shape_array.flatten().astype('int64')))
+        else:
+            raise "Invalid dtype"
         expand_node = helper.make_node("Expand", ["in", "shape"], ["out"])
 
         graph = helper.make_graph([shape_node, expand_node],
@@ -196,13 +207,15 @@ def _test_expand(name, data, shape, ref_data):
     shape = (3, 4)
     data = np.random.uniform(size=in_shape).astype(np.float32)
     ref_data = np.tile(data, 4)
-    _test_expand('expand_with_dim_unchanged_test', data, shape, ref_data)
+    _test_expand('expand_with_dim_unchanged_test', data, shape, ref_data, "int32")
+    _test_expand('expand_with_dim_unchanged_test', data, shape, ref_data, "int64")
 
     in_shape = (3, 1)
     shape = (2, 1, 6)
     data = np.random.uniform(size=in_shape).astype(np.float32)
     ref_data = data * np.ones(shape, dtype=np.float32)
-    _test_expand('expand_with_dim_changed_test', data, shape, ref_data)
+    _test_expand('expand_with_dim_changed_test', data, shape, ref_data, "int32")
+    _test_expand('expand_with_dim_changed_test', data, shape, ref_data, "int64")
 
 
 def verify_depth_to_space(inshape, outshape, mode, blockSize):
@@ -822,8 +835,8 @@ def verify_batch_matmul(a_shape, b_shape):
             model, [a_array, b_array], target, ctx)
         tvm.testing.assert_allclose(out_np, tvm_out, rtol=1e-5, atol=1e-5)
 
-# TODO(mbrookhart): enable once VM supports heterogenous execution
-# @tvm.testing.uses_gpu
+# TODO(mbrookhart): enable cuda once VM supports heterogenous execution
+@tvm.testing.parametrize_targets("llvm")
 def test_batch_matmul():
     verify_batch_matmul((2, 3, 4, 3), (2, 3, 3, 4))
     verify_batch_matmul((2, 4, 3), (3, 4))
@@ -1024,11 +1037,9 @@ def _test_upsample_bilinear_opset9():
         graph, producer_name='upsample_bilinear_opset9_test')
 
     for target, ctx in tvm.testing.enabled_targets():
-        tvm_out = get_tvm_output(
-            model, in_array, target, ctx, out_shape, 'float32')
+        tvm_out = get_tvm_output_with_vm(model, [in_array], target, ctx, opset=9, freeze_params=True)
         tvm.testing.assert_allclose(out_array, tvm_out, rtol=1e-5, atol=1e-5)
 
-
 def _test_upsample3d_trilinear():
     scale = 2
     in_shape = (1, 1, 3, 3, 3)
@@ -1062,7 +1073,8 @@ def _test_upsample3d_trilinear():
             model, in_array, target, ctx, out_shape, 'float32')
         tvm.testing.assert_allclose(out_array, tvm_out, rtol=1e-5, atol=1e-5)
 
-@tvm.testing.uses_gpu
+# TODO(mbrookhart): enable once VM supports heterogenous execution
+# @tvm.testing.uses_gpu
 def test_upsample():
     _test_upsample_nearest()
     _test_upsample_bilinear()
@@ -1455,7 +1467,7 @@ def verify_pad_v11(indata, pads, mode='constant', value=0.0):
                                   outputs=[helper.make_tensor_value_info("output",
                                                                          TensorProto.FLOAT, list(outdata.shape))])
     else:
-        inputs = [indata, pads, np.array([value])]
+        inputs = [indata, pads, np.array([value]).astype("float32")]
         outdata = np.pad(indata, pad_width=np_pads,
                          mode='constant', constant_values=value)
         node = helper.make_node(
@@ -1471,7 +1483,7 @@ def verify_pad_v11(indata, pads, mode='constant', value=0.0):
                                           helper.make_tensor_value_info("pads",
                                                                         TensorProto.INT64,(len(pads),)),
                                           helper.make_tensor_value_info("constant_value",
-                                                                        TensorProto.INT64,(1,)),
+                                                                        TensorProto.FLOAT,(1,)),
                                           ],
                                   initializer=[helper.make_tensor("pads", TensorProto.INT64, (len(pads),), pads),
                                                helper.make_tensor("constant_value", TensorProto.FLOAT, (1,), [value])],
@@ -1480,12 +1492,12 @@ def verify_pad_v11(indata, pads, mode='constant', value=0.0):
     model = helper.make_model(graph, producer_name='pad_test')
     #  tvm result
     for target, ctx in tvm.testing.enabled_targets():
-        tvm_out = get_tvm_output(
-            model, inputs, target, ctx, outdata.shape, 'float32', opset=11)
+        tvm_out = get_tvm_output_with_vm(model, inputs, target, ctx, opset=11, freeze_params=False)
     tvm.testing.assert_allclose(outdata, tvm_out, rtol=1e-5, atol=1e-5)
 
 
-@tvm.testing.uses_gpu
+# TODO(mbrookhart): enable once VM supports heterogenous execution
+# @tvm.testing.uses_gpu
 def test_pad():
     verify_pad(np.random.randn(2, 2).astype(
         np.float32), [0, 1, 0, 0], 'constant', 0.0)

tests/python/relay/test_op_level10.py

@@ -358,7 +358,7 @@ def verify_dynamic_batch_matmul(x_shape, y_shape, out_shape, dtype="float32"):
     y_np = np.random.uniform(size=y_shape).astype(dtype)
     z_np = tvm.topi.testing.batch_matmul(x_np, y_np)
 
-    for target, ctx in ctx_list():
+    for target, ctx in tvm.testing.enabled_targets():
         for kind in ["vm", "debug"]:
             mod = tvm.ir.IRModule.from_expr(func)
             intrp = relay.create_executor(kind, mod=mod, ctx=ctx, target=target)