Add zeropad layers (#13)

Author: Laicheng0830

tests/test_padding.py

@@ -14,9 +14,11 @@ class CustomModel(Module):
     def __init__(self):
         super(CustomModel, self).__init__(name="custom")
         self.pad = PadLayer([[1, 2], [3, 4], [5, 6], [7, 8]], "REFLECT", name='inpad')
+        self.pad2d = ZeroPad2d(padding=((2, 2), (3, 3)), data_format='channels_last')
 
     def forward(self, inputs):
         x = self.pad(inputs)
+        x = self.pad2d(x)
         return x
 
 net = CustomModel()

tlx2onnx/op_mapper/nn/padding.py

@@ -22,7 +22,7 @@ def version_1(cls, node, **kwargs):
         layer = node['node'].layer
         # get attrs
         value = np.array(layer.constant_values).astype(node['dtype'])
-        c_value = numpy_helper.from_array(value, name='value')
+        c_value = numpy_helper.from_array(value, name=layer.name + 'value')
         onnx_init.append(c_value)
         # processing mode
         mode_dict = {"CONSTANT": 'constant', "REFLECT": 'reflect',"SYMMETRIC": 'edge'}
@@ -39,17 +39,71 @@ def version_1(cls, node, **kwargs):
         for i in range(len(pads_temp) // 2):
             pads.append(pads_temp[i*2+1])
         pads = np.array(pads).astype(np.int64)
-        p_value = numpy_helper.from_array(pads, name='pads')
+        p_value = numpy_helper.from_array(pads, name=layer.name + 'pads')
         onnx_init.append(p_value)
         # make nodes
         v_out = helper.make_tensor_value_info(out_name, dtype, shape=out_shape)
         onnx_value.append(v_out)
 
         if mode == 'constant':
-            p_node, out = make_node('Pad', inputs=[in_name, 'pads', 'value'], outputs=[out_name], mode='constan')
+            p_node, out = make_node('Pad', inputs=[in_name, layer.name + 'pads', layer.name + 'value'], outputs=[out_name], mode='constant')
             onnx_node.append(p_node)
         else:
-            p_node, out = make_node('Pad', inputs=[in_name, 'pads'], outputs=[out_name], mode=mode)
+            p_node, out = make_node('Pad', inputs=[in_name, layer.name + 'pads'], outputs=[out_name], mode=mode)
             onnx_node.append(p_node)
 
-        return onnx_node, onnx_value, onnx_init
+        return onnx_node, onnx_value, onnx_init
+
+
+@OpMapper(['ZeroPad1d', 'ZeroPad2d', 'ZeroPad3d'])
+class ZeroPad():
+    # supports v1-v12
+
+    @classmethod
+    def version_1(cls, node, **kwargs):
+        onnx_node, onnx_value, onnx_init = [], [], []
+        # get inputs outputs
+        in_name = node['in_nodes_name'][0]
+        out_name = node['out_nodes_name'][0]
+        out_shape = node['out_tensors'][0]
+        dtype = NP_TYPE_TO_TENSOR_TYPE[node['dtype']]
+        layer = node['node'].layer
+        # get attrs
+        padding = layer.padding
+        data_format = layer.data_format
+        pads_temp = convert_padding(padding, data_format)
+
+        pads = []
+        for i in range(len(pads_temp)//2):
+            pads.append(pads_temp[2*i])
+        for i in range(len(pads_temp) // 2):
+            pads.append(pads_temp[i*2+1])
+        pads = np.array(pads).astype(np.int64)
+
+        p_value = numpy_helper.from_array(pads, name=layer.name + 'pads')
+        onnx_init.append(p_value)
+
+        # make nodes
+        v_out = helper.make_tensor_value_info(out_name, dtype, shape=out_shape)
+        onnx_value.append(v_out)
+        p_node, out = make_node('Pad', inputs=[in_name, layer.name + 'pads'], outputs=[out_name], mode='constant')
+        onnx_node.append(p_node)
+
+        return onnx_node, onnx_value, onnx_init
+
+
+def convert_padding(padding, data_format):
+    if np.size(padding) == 2:
+        if data_format == 'channels_first':
+            out = (0, 0, 0, 0) + padding
+        else:
+            out = (0, 0) + padding + (0, 0)
+    else:
+        pads_temp = padding[0]
+        for i in np.arange(1, len(padding)):
+            pads_temp += padding[i]
+        if data_format == 'channels_first':
+            out = (0, 0, 0, 0) + pads_temp
+        else:
+            out = (0, 0) + pads_temp + (0, 0)
+    return out