diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index a8cb216e26bd..36ef2bdb67b5 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -2759,6 +2759,38 @@ def _impl_v11(cls, bb, inputs, attr, params): # edge mode - replicate border values return bb.emit_te(topi.nn.replicate_pad, inputs[0], pad_before, pad_after) + @classmethod + def _impl_v19(cls, bb, inputs, attr, params): + pads = get_constant(inputs[1], params) + constant_value = get_constant(inputs[2], params) + if constant_value is not None: + constant_value = constant_value.data.numpy().item() + else: + constant_value = 0.0 + + if isinstance(pads, relax.Constant): + pad_before, pad_after = _np.split(pads.data.numpy(), 2) + pad_before = _np.ndarray.tolist(pad_before) + pad_after = _np.ndarray.tolist(pad_after) + else: + raise ValueError("Dynamic pads are not supported yet.") + + pad_mode = attr.get("mode", b"constant").decode("utf-8") + if pad_mode not in ["constant", "edge", "reflect", "wrap"]: + raise tvm.error.OpAttributeInvalid( + "Value " + pad_mode + ' in attribute "mode" is invalid for operator Pad.' + ) + + if pad_mode == "constant": + return bb.emit_te(topi.nn.pad, inputs[0], pad_before, pad_after, constant_value) + elif pad_mode == "reflect": + return bb.emit_te(topi.nn.mirror_pad, inputs[0], pad_before, pad_after, "REFLECT") + elif pad_mode == "wrap": + return bb.emit_te(topi.nn.circular_pad, inputs[0], pad_before, pad_after) + else: + # edge mode - replicate border values + return bb.emit_te(topi.nn.replicate_pad, inputs[0], pad_before, pad_after) + class Tile(OnnxOpConverter): """Converts an onnx Tile node into an equivalent Relax expression.""" diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 5aff95da5a46..6630e0f16c60 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -3656,14 +3656,14 @@ def test_pad(dynamic): if dynamic: pytest.skip("Dynamic pad not supported") - def verify_pad(input_shape, pads, mode="constant", value=0.0): + def verify_pad(input_shape, pads, mode="constant", value=0.0, opset=14): indata = np.random.normal(size=input_shape).astype(np.float32) # numpy expect result len_dim = len(pads) // 2 np_pads = [(pads[i], pads[i + len_dim]) for i in range(len_dim)] pads = np.array(pads) # onnx graph - if mode in ["edge", "reflect"]: + if mode in ["edge", "reflect", "wrap"]: outdata = np.pad(indata, pad_width=np_pads, mode=mode) node = helper.make_node("Pad", inputs=["input", "pads"], outputs=["output"], mode=mode) graph = helper.make_graph( @@ -3700,7 +3700,7 @@ def verify_pad(input_shape, pads, mode="constant", value=0.0): ], ) model = helper.make_model(graph, producer_name="pad_test") - check_correctness(model) + check_correctness(model, opset=opset) verify_pad((2, 2), [0, 1, 0, 0], "constant", 0.0) verify_pad((2, 3), [1, 0, 0, 1], "constant", 0.0) @@ -3708,6 +3708,7 @@ def verify_pad(input_shape, pads, mode="constant", value=0.0): verify_pad((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "reflect") verify_pad((2, 3), [1, 1, 1, 1], "edge") verify_pad((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "edge") + verify_pad((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "wrap", opset=19) @pytest.mark.parametrize("dynamic", [True, False])