diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 5471f67ea106..6e05646b935a 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -961,6 +961,18 @@ def _impl_v1(cls, inputs, attr, params): return _op.power(out, reci_p) +class GlobalLpPool(OnnxOpConverter): + """Operator converter for GlobalLpPool.""" + + @classmethod + def _impl_v1(cls, inputs, attr, params): + # TODO: GlobalLpPool does not yet support dynamic shapes + in_shape = infer_shape(inputs[0]) + attr["kernel_shape"] = in_shape[2:] + + return LpPool._impl_v1(inputs, attr, params) + + class Mul(Elemwise): """Operator converter for Multiply.""" @@ -3578,6 +3590,7 @@ def _get_convert_map(opset): # defs/nn "AveragePool": AveragePool.get_converter(opset), "LpPool": LpPool.get_converter(opset), + "GlobalLpPool": GlobalLpPool.get_converter(opset), "MaxPool": MaxPool.get_converter(opset), "MaxUnpool": MaxUnpool.get_converter(opset), "Conv": Conv.get_converter(opset), diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index 9e0eb1f75217..e4cb83a5cb05 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -3298,6 +3298,49 @@ def verify_lppool(x_shape, kernel_shape, p, strides, pads, out_shape, auto_pad=" ) +def verify_global_lppool(x_shape, p, out_shape, target, dev): + pool_node = helper.make_node( + "GlobalLpPool", + inputs=["x"], + outputs=["y"], + p=p, + ) + + graph = helper.make_graph( + [pool_node], + "global_lppool_test", + inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, list(x_shape))], + outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, list(out_shape))], + ) + + model = helper.make_model(graph, producer_name="global_lppool_test") + verify_with_ort( + model, [x_shape], out_shape, use_vm=True, convert_to_static=True, target=target, dev=dev + ) + + +@tvm.testing.parametrize_targets +def test_global_lppool(target, dev): + + # LpPool1D + verify_global_lppool(x_shape=[1, 15, 16], p=2, out_shape=[1, 15, 1], target=target, dev=dev) + + # LpPool2D + verify_global_lppool( + x_shape=[1, 15, 32, 32], p=2, out_shape=[1, 15, 1, 1], target=target, dev=dev + ) + + # LpPool2D + verify_global_lppool( + x_shape=[1, 15, 32, 32], p=3, out_shape=[1, 15, 1, 1], target=target, dev=dev + ) + + # LpPool3D + verify_global_lppool( + x_shape=[1, 15, 3, 32, 32], p=2, out_shape=[1, 15, 1, 1, 1], target=target, dev=dev + ) + + def verify_rnn( seq_length, batch_size, @@ -5644,3 +5687,4 @@ def repeat(N, D): test_random_uniform() test_convinteger() test_batch_matmul() + test_global_lppool()