diff --git a/python/tvm/contrib/target/onnx.py b/python/tvm/contrib/target/onnx.py index b05265fa976a..6189e6ccf6bc 100644 --- a/python/tvm/contrib/target/onnx.py +++ b/python/tvm/contrib/target/onnx.py @@ -147,6 +147,7 @@ def convert_attributes(cls, attrs): "pads": attrs.get_int_tuple("padding"), "strides": attrs.get_int_tuple("strides"), "kernel_shape": attrs.get_int_tuple("pool_size"), + "ceil_mode": 1 if attrs.ceil_mode else 0, } @@ -330,7 +331,10 @@ def convert_attributes(cls, attrs): after.append(axis_pads[1]) pads = before + after pads = numpy.asarray(pads, dtype=pads[0].dtype) - return {"pads": pads, "mode": attrs.get_str("pad_mode"), "constant_value": attrs.pad_value} + return { + "pads": pads, + "mode": attrs.get_str("pad_mode"), + } @classmethod def convert(cls, node_entry, model_container, node_dict): @@ -341,16 +345,17 @@ def convert(cls, node_entry, model_container, node_dict): attrs = cls.convert_attributes(node_entry["relay_node"].attrs) name = node_entry["name"] - data = numpy.asarray(attrs["pads"], dtype=attrs["pads"][0].dtype).astype(numpy.int64) - value = numpy.dtype(node_entry["types"][0].dtype).type(attrs["constant_value"]) + pad_data = numpy.asarray(attrs["pads"], dtype=attrs["pads"][0].dtype).astype(numpy.int64) input_names = [ node_entry["input_names"][0], - add_input(data, name, "pads", model_container), - add_input(value, name, "value", model_container), + add_input(pad_data, name, "pads", model_container), + node_entry["input_names"][1], ] - node = onnx.helper.make_node(cls.__name__, input_names, node_entry["output_names"]) + node = onnx.helper.make_node( + cls.__name__, input_names, node_entry["output_names"], mode=attrs["mode"] + ) model_container.add_nodes([node]) diff --git a/tests/python/contrib/test_onnx.py b/tests/python/contrib/test_onnx.py index 3636409f8a06..7d6a2cab5bc8 100644 --- a/tests/python/contrib/test_onnx.py +++ b/tests/python/contrib/test_onnx.py @@ -27,6 +27,7 @@ import tvm from tvm import relay from tvm.contrib.target.onnx import to_onnx +from tvm.relay.testing import run_infer_type def func_to_onnx(func, name): @@ -270,14 +271,16 @@ def verify_batch_norm(axis=1): def test_pad(): + """Pad unit test.""" + def verify_pad(): - for dtype in ["float16", "float32"]: - dshape = (4, 10, 7, 7) - x = relay.var("x", shape=dshape, dtype=dtype) - y = relay.nn.pad(x, ((1, 1), (2, 2), (3, 3), (4, 4))) - func = relay.Function([x], y) - x_data = np.random.uniform(size=dshape).astype(dtype) - verify_results(func, [x_data], "test_pad", rtol=1e-5, atol=1e-5) + dshape = (4, 10, 7, 7) + x = relay.var("x", shape=dshape, dtype="int32") + y = relay.nn.pad(x, ((1, 1), (2, 2), (3, 3), (4, 4))) + func = relay.Function([x], y) + func = run_infer_type(func) + x_data = np.random.randint(low=-255, high=255, size=dshape).astype(np.int32) + verify_results(func, [x_data], "test_pad", rtol=1e-5, atol=1e-5) verify_pad()