Skip to content

Commit

Permalink
Frontend: add onnx GlobalLpPool op
Browse files Browse the repository at this point in the history
  • Loading branch information
xp224797 authored and 徐鹏 committed Aug 25, 2021
1 parent b819364 commit 4258567
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 0 deletions.
13 changes: 13 additions & 0 deletions python/tvm/relay/frontend/onnx.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
in_shape = infer_shape(inputs[0])
assert len(in_shape) >= 4, "input shape support NCHW or NCDHW"
attr["kernel_shape"] = in_shape[2:]

return LpPool._impl_v1(inputs, attr, params)


class Mul(Elemwise):
"""Operator converter for Multiply."""

Expand Down Expand Up @@ -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),
Expand Down
45 changes: 45 additions & 0 deletions tests/python/frontend/onnx/test_forward.py
Original file line number Diff line number Diff line change
Expand Up @@ -3298,6 +3298,50 @@ def verify_lppool(x_shape, kernel_shape, p, strides, pads, out_shape, auto_pad="
)


def verify_global_lppool(x_shape, p, out_shape):
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)


@tvm.testing.uses_gpu
def test_global_lppool():

# Pool2D
verify_global_lppool(
x_shape=[1, 15, 32, 32],
p=2,
out_shape=[1, 15, 1, 1],
)

# Pool2D
verify_global_lppool(
x_shape=[1, 15, 32, 32],
p=3,
out_shape=[1, 15, 1, 1],
)

# Pool3D
verify_global_lppool(
x_shape=[1, 15, 3, 32, 32],
p=2,
out_shape=[1, 15, 1, 1, 1],
)


def verify_rnn(
seq_length,
batch_size,
Expand Down Expand Up @@ -5644,3 +5688,4 @@ def repeat(N, D):
test_random_uniform()
test_convinteger()
test_batch_matmul()
test_global_lppool()

0 comments on commit 4258567

Please sign in to comment.