From d391753839065fb3cd1ed8a818699634e5e69a0e Mon Sep 17 00:00:00 2001 From: Samuel Date: Sat, 4 Apr 2020 15:03:43 +0530 Subject: [PATCH] [ONNX]Pool3d & upsample3d op support (#5135) * [ONNX]Pool3d and Upsample3d op updated * Pool3d and Upsample3d testcase * Review comments fixed * Review comments --- python/tvm/relay/frontend/onnx.py | 32 ++++++--- tests/python/frontend/onnx/test_forward.py | 77 +++++++++++++++++++++- 2 files changed, 100 insertions(+), 9 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index beb8e8559dab..527a1ed2f07b 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -137,8 +137,10 @@ def onnx_default_layout(dims): return 'NCW' if dims == 2: return 'NCHW' + if dims == 3: + return 'NCDHW' - msg = "Only 1d and 2d layouts are currently supported" + msg = "Only 1D, 2D and 3D layouts are currently supported" raise tvm.error.OpAttributeInvalid(msg.format(op_name)) @@ -151,8 +153,10 @@ def onnx_storage_order2layout(storage_order, dims=2): return 'NCW' if storage_order == 0 else 'NWC' if dims == 2: return 'NCHW' if storage_order == 0 else 'NHWC' + if dims == 3: + return 'NCDHW' if storage_order == 0 else 'NDHWC' - msg = "Only 1d and 2d layouts are currently supported" + msg = "Only 1D, 2D and 3D layouts are currently supported" raise tvm.error.OpAttributeInvalid(msg.format(op_name)) @@ -780,19 +784,31 @@ def _impl_v9(cls, inputs, attr, params): assert len(inputs) == 2, "Upsample op take 2 inputs, {} given".format(len(inputs)) scales = params[inputs[1].name_hint].asnumpy() inputs = inputs[:1] - assert len(scales) == 4 and scales[0] == 1.0 and scales[1] == 1.0 + assert scales[0] == 1.0 and scales[1] == 1.0 + input_shape = infer_shape(inputs[0]) + dims = len(input_shape) mode = attr.get('mode') if mode == b'nearest': method = "nearest_neighbor" elif mode == b'linear': - method = "bilinear" + method = "trilinear" if dims == 5 else "bilinear" 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, - 'layout': 'NCHW', 'align_corners': True} - return AttrCvt('upsampling')(inputs, attr) - + attr = {'scale_h': scales[-2], + 'scale_w': scales[-1], + 'method': method} + if dims == 5: + assert len(scales) == 5 + attr['scale_d'] = scales[-3] + attr['layout'] = 'NCDHW' + op_name = 'upsampling3d' + else: + assert len(scales) == 4 + attr['layout'] = 'NCHW' + attr['align_corners'] = True + op_name = 'upsampling' + return AttrCvt(op_name)(inputs, attr) class Shape(OnnxOpConverter): """ Operator converter for Shape. diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index 917ec991351b..2c0849451a25 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -741,6 +741,30 @@ def _test_upsample_nearest(): tvm.testing.assert_allclose(out_array, tvm_out) +def _test_upsample3d_nearest(): + scale = 2 + in_shape = (1, 1, 3, 3, 3) + out_shape = (1, 1, 3*scale, 3*scale, 3*scale) + y = helper.make_node("Upsample", ['in'], [ + 'out'], mode='nearest', scales=[1.0, 1.0, 2.0, 2.0, 2.0]) + + in_array = np.random.uniform(size=in_shape).astype(np.float32) + out_array = topi.testing.upsampling3d_python( + in_array, (scale, scale, scale), "NCDHW") + + graph = helper.make_graph([y], + 'upsample_nearest_test', + inputs=[helper.make_tensor_value_info( + "in", TensorProto.FLOAT, list(in_shape))], + outputs=[helper.make_tensor_value_info("out", TensorProto.FLOAT, list(out_shape))]) + + model = helper.make_model(graph, producer_name='upsample_nearest_test') + + for target, ctx in ctx_list(): + tvm_out = get_tvm_output( + model, in_array, target, ctx, out_shape, 'float32') + tvm.testing.assert_allclose(out_array, tvm_out) + def _test_upsample_bilinear(): scale = 2 in_shape = (1, 1, 3, 3) @@ -800,11 +824,45 @@ def _test_upsample_bilinear_opset9(): 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) + out_shape = (1, 1, 3*scale, 3*scale, 3*scale) + y = helper.make_node("Upsample", ['in', 'scales'], ['out'], mode='linear') + scales = [1.0, 1.0, 2.0, 2.0, 2.0] + in_array = np.random.uniform(size=in_shape).astype(np.float32) + out_array = topi.testing.trilinear_resize3d_python( + in_array, (3*scale, 3*scale, 3*scale), "NCDHW", coordinate_transformation_mode="half_pixel") + + ref_array = np.array(scales) + ref_node = helper.make_node('Constant', + inputs=[], + outputs=['scales'], + value=onnx.helper.make_tensor(name='const_tensor', + data_type=TensorProto.FLOAT, + dims=ref_array.shape, + vals=ref_array.flatten().astype(float))) + + graph = helper.make_graph([ref_node, y], + 'upsample_trilinear_test', + inputs=[helper.make_tensor_value_info( + "in", TensorProto.FLOAT, list(in_shape))], + outputs=[helper.make_tensor_value_info("out", TensorProto.FLOAT, list(out_shape))]) + + model = helper.make_model( + graph, producer_name='upsample_trilinear_test') + + for target, ctx in ctx_list(): + tvm_out = get_tvm_output( + model, in_array, target, ctx, out_shape, 'float32') + tvm.testing.assert_allclose(out_array, tvm_out, rtol=1e-5, atol=1e-5) + def test_upsample(): _test_upsample_nearest() _test_upsample_bilinear() _test_upsample_bilinear_opset9() - + _test_upsample3d_nearest() + _test_upsample3d_trilinear() def _test_softmax(inshape, axis): opname = 'Softmax' @@ -1999,6 +2057,23 @@ def test_pooling(): mode=mode, auto_pad='SAME_UPPER') + # Pool3D with stride + verify_pooling(x_shape=[1, 1, 32, 32, 32], + kernel_shape=[3, 3, 3], + strides=[2, 2, 2], + pads=[1, 1, 1, 1, 1, 1], + out_shape=[1, 1, 16, 16, 16], + mode=mode) + + # Pool3D with stride and autopadding + verify_pooling(x_shape=[1, 1, 32, 32, 32], + kernel_shape=[3, 3, 3], + strides=[2, 2, 2], + pads=None, + out_shape=[1, 1, 16, 16, 16], + mode=mode, + auto_pad='SAME_UPPER') + def verify_lstm(seq_length, batch_size,