Unverified Commit fd9ce583 by Samuel Committed by GitHub

[ONNX]Pool3d & upsample3d op support (#5135)

* [ONNX]Pool3d and Upsample3d op updated

* Pool3d and Upsample3d testcase

* Review comments fixed

* Review comments
parent 0cfdecda
......@@ -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 @@ class Upsample(OnnxOpConverter):
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.
......
......@@ -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,
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment