Add Onnx Pad v11 (#5539)

diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py
index 4ae083c..f4bcb6b 100644 (file)
--- a/python/tvm/relay/frontend/onnx.py
+++ b/python/tvm/relay/frontend/onnx.py
@@ -557,6 +557,31 @@ class Pad(OnnxOpConverter):
             },
             )(inputs, attr, params)
 
+    @classmethod
+    def _impl_v11(cls, inputs, attr, params):
+        pad_width = []
+        pads = infer_value_simulated(inputs[1], params).asnumpy()
+        if len(inputs) == 3:
+            value = infer_value_simulated(inputs[2], params).asnumpy().item()
+        else:
+            value = 0
+        attr["pad_value"] = value
+        dims = int(len(pads) / 2)
+        for i in range(dims):
+            pad_width.append((pads[i], pads[i+dims]))
+        attr['pad_width'] = pad_width
+        pad_mode = attr.get('mode', b'constant').decode('utf-8')
+        if pad_mode in ['constant', 'edge', 'reflect']:
+            attr['pad_mode'] = pad_mode
+            attr.pop('mode', None)
+        else:
+            raise tvm.error.OpAttributeInvalid(
+                'Value ' + pad_mode + ' in attribute "mode" is invalid for operator Pad.')
+
+        return AttrCvt('pad')(inputs[:1], attr, params)
+
+
+
 
 class ParametricSoftPlus(OnnxOpConverter):
     """ Operator converter for ParametricSoftPlus.

diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py
index 78658e7..8bc9f45 100644 (file)
--- a/tests/python/frontend/onnx/test_forward.py
+++ b/tests/python/frontend/onnx/test_forward.py
@@ -1279,7 +1279,63 @@ def verify_pad(indata, pads, mode='constant', value=0.0):
     #  tvm result
     for target, ctx in ctx_list():
         tvm_out = get_tvm_output(
-            model, indata, target, ctx, outdata.shape, 'float32')
+            model, indata, target, ctx, outdata.shape, 'float32', opset=2)
+    tvm.testing.assert_allclose(outdata, tvm_out, rtol=1e-5, atol=1e-5)
+
+
+def verify_pad_v11(indata, pads, mode='constant', value=0.0):
+    indata = np.array(indata).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']:
+        inputs = [indata, pads]
+        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([node],
+                                  'pad_test',
+                                  inputs=[helper.make_tensor_value_info("input",
+                                                                        TensorProto.FLOAT, list(indata.shape)),
+                                          helper.make_tensor_value_info("pads",
+                                                                        TensorProto.INT64,(len(pads),))],
+                                  initializer=[helper.make_tensor("pads", TensorProto.INT64, (len(pads),), pads)],
+                                  outputs=[helper.make_tensor_value_info("output",
+                                                                         TensorProto.FLOAT, list(outdata.shape))])
+    else:
+        inputs = [indata, pads, np.array([value])]
+        outdata = np.pad(indata, pad_width=np_pads,
+                         mode='constant', constant_values=value)
+        node = helper.make_node(
+            'Pad',
+            inputs=['input', 'pads', 'constant_value'],
+            outputs=['output'],
+            mode='constant'
+        )
+        graph = helper.make_graph([node],
+                                  'pad_test',
+                                  inputs=[helper.make_tensor_value_info("input",
+                                                                        TensorProto.FLOAT, list(indata.shape)),
+                                          helper.make_tensor_value_info("pads",
+                                                                        TensorProto.INT64,(len(pads),)),
+                                          helper.make_tensor_value_info("constant_value",
+                                                                        TensorProto.INT64,(1,)),
+                                          ],
+                                  initializer=[helper.make_tensor("pads", TensorProto.INT64, (len(pads),), pads),
+                                               helper.make_tensor("constant_value", TensorProto.FLOAT, (1,), [value])],
+                                  outputs=[helper.make_tensor_value_info("output",
+                                                                         TensorProto.FLOAT, list(outdata.shape))])
+    model = helper.make_model(graph, producer_name='pad_test')
+    #  tvm result
+    for target, ctx in ctx_list():
+        tvm_out = get_tvm_output(
+            model, inputs, target, ctx, outdata.shape, 'float32', opset=11)
     tvm.testing.assert_allclose(outdata, tvm_out, rtol=1e-5, atol=1e-5)
 
 
@@ -1295,6 +1351,17 @@ def test_pad():
     verify_pad(np.random.randn(1, 3, 4, 5).astype(
         np.float32), [0, 0, 1, 1, 0, 0, 1, 1], 'reflect')
 
+    verify_pad_v11(np.random.randn(2, 2).astype(
+        np.float32), [0, 1, 0, 0], 'constant', 0.0)
+    verify_pad_v11(np.random.randn(2, 3).astype(
+        np.float32), [1, 0, 0, 1], 'constant', 0.0)
+    verify_pad_v11(np.random.randn(3, 2).astype(
+        np.float32), [0, 0, 1, 0], 'constant', 5.0)
+    verify_pad_v11(np.random.randn(1, 3, 4, 5).astype(
+        np.float32), [0, 0, 1, 1, 0, 0, 1, 1], 'edge')
+    verify_pad_v11(np.random.randn(1, 3, 4, 5).astype(
+        np.float32), [0, 0, 1, 1, 0, 0, 1, 1], 'reflect')
+
 
 def verify_reduce_x(name, indata, axis, keepdims):
     indata = np.array(indata).astype(np.float32)