[im2col, test_im2col]: added support for non-equal padding

src/finn/custom_op/general/im2col.py

@@ -10,8 +10,19 @@
 # of shape (batches, channels, height, width)
 
 
+def compute_conv_output_dim_2D_padding(ifm_dim, k, stride, pad=0):
+    """Returns spatial output dimension size for convolution with given params.
+    Pad denotes the total amount of padding among a certain dimension"""
+    if ifm_dim == 1:
+        out_dim = 1
+    else:
+        out_dim = int(((ifm_dim + pad - k) / stride) + 1)
+    return out_dim
+
+
 def compute_conv_output_dim(ifm_dim, k, stride, pad=0):
-    """Returns spatial output dimension size for convolution with given params."""
+    """Returns spatial output dimension size for convolution with given params.
+    Pad is either 0 or 1"""
     if ifm_dim == 1:
         out_dim = 1
     else:
@@ -25,8 +36,10 @@ def get_im2col_indices_nchw(
     """Returns im2col indices."""
     # First figure out what the size of the output should be
     N, C, H, W = x_shape
-    out_height = compute_conv_output_dim(H, field_height, stride_y, padding)
-    out_width = compute_conv_output_dim(W, field_width, stride_x, padding)
+    pad_H = padding[0] + padding[2]
+    pad_W = padding[1] + padding[3]
+    out_height = compute_conv_output_dim_2D_padding(H, field_height, stride_y, pad_H)
+    out_width = compute_conv_output_dim_2D_padding(W, field_width, stride_x, pad_W)
 
     i0 = np.repeat(np.arange(field_height), field_width)
     i0 = np.tile(i0, C)
@@ -53,21 +66,21 @@ def im2col_indices_nchw(
     if H == 1:  # Shape of input image is: (1, C, 1, W)
         x_padded = np.pad(
             x,
-            ((0, 0), (0, 0), (0, 0), (p, p)),
+            ((0, 0), (0, 0), (0, 0), (p[1], p[3])),
             mode="constant",
             constant_values=pad_val,
         )
     elif W == 1:  # Shape of input image is: (1, C, H, 1)
         x_padded = np.pad(
             x,
-            ((0, 0), (0, 0), (p, p), (0, 0)),
+            ((0, 0), (0, 0), (p[0], p[2]), (0, 0)),
             mode="constant",
             constant_values=pad_val,
         )
     elif H > 1 and W > 1:  # Shape of input image is: (1, C, H, W)
         x_padded = np.pad(
             x,
-            ((0, 0), (0, 0), (p, p), (p, p)),
+            ((0, 0), (0, 0), (p[0], p[2]), (p[1], p[3])),
             mode="constant",
             constant_values=pad_val,
         )
@@ -101,7 +114,7 @@ def get_nodeattr_types(self):
             "stride": ("i", True, 1),
             "kernel_size": ("ints", True, []),
             "input_shape": ("s", True, ""),
-            "pad_amount": ("i", False, 0),
+            "pad_amount": ("ints", False, [0, 0, 0, 0]),  # default: no padding
             "pad_value": ("i", False, 0),
             # depthwise: if 1, infer ConvolutionInputGenerator with depthwise == 1
             "depthwise": ("i", False, 0, {0, 1}),
@@ -111,7 +124,11 @@ def make_shape_compatible_op(self, model):
         k = self.get_nodeattr("kernel_size")  # Assumption: Height x Width
         stride = self.get_nodeattr("stride")
         ishape = self.get_nodeattr("input_shape")
-        pad = self.get_nodeattr("pad_amount")
+        pad = self.get_nodeattr(
+            "pad_amount"
+        )  # padding: [H_begin, W_begin, H_end, W_end]
+        pad_H = pad[0] + pad[2]
+        pad_W = pad[1] + pad[3]
 
         # convert string into list of integers
         ishape = ishape.strip("(")
@@ -142,8 +159,8 @@ def make_shape_compatible_op(self, model):
                 k_W == 1
             ), "Unexpected kernel shape for input image of dimensions (N, H, 1, C)"
 
-        ofm_dim_H = compute_conv_output_dim(ifm_dim_H, k_H, stride, pad)
-        ofm_dim_W = compute_conv_output_dim(ifm_dim_W, k_W, stride, pad)
+        ofm_dim_H = compute_conv_output_dim_2D_padding(ifm_dim_H, k_H, stride, pad_H)
+        ofm_dim_W = compute_conv_output_dim_2D_padding(ifm_dim_W, k_W, stride, pad_W)
 
         # implement tensor with correct shape
         values = np.random.randn(1, ofm_dim_H, ofm_dim_W, k_H * k_W * ifm_ch).astype(
@@ -180,15 +197,18 @@ def execute_node(self, context, graph):
 
         stride = self.get_nodeattr("stride")
         pad = self.get_nodeattr("pad_amount")
+        pad_H = pad[0] + pad[2]
+        pad_W = pad[1] + pad[3]
         pad_val = self.get_nodeattr("pad_value")
         iname = node.input[0]
         x = context[iname]
         qnt_annotations = graph.quantization_annotation
         ret = util.get_by_name(qnt_annotations, iname, "tensor_name")
         ret = util.get_by_name(ret.quant_parameter_tensor_names, "finn_datatype", "key")
         idt = DataType[ret.value]
-        if pad != 0:
-            assert idt.allowed(pad_val), "Im2Col dtype must allow pad_val"
+        for val in pad:
+            if val != 0:
+                assert idt.allowed(val), "Im2Col dtype must allow pad_val"
         # check that input is NHWC
         assert x.ndim == 4, "Unexpected number of input dims for Im2Col"
         N, H, W, C = x.shape
@@ -202,8 +222,9 @@ def execute_node(self, context, graph):
                 k_W == 1
             ), "Unexpected kernel shape for input image of dimensions (N, H, 1, C)"
 
-        out_dim_H = compute_conv_output_dim(H, k_H, stride, pad)
-        out_dim_W = compute_conv_output_dim(W, k_W, stride, pad)
+        out_dim_H = compute_conv_output_dim_2D_padding(H, k_H, stride, pad_H)
+        out_dim_W = compute_conv_output_dim_2D_padding(W, k_W, stride, pad_W)
+
         # internally convert input to NCHW
         x = x.transpose(0, 3, 1, 2)
         # call NCHW im2col implementation

tests/custom_op/test_im2col.py

@@ -4,7 +4,7 @@
 import finn.core.onnx_exec as oxe
 from finn.core.datatype import DataType
 from finn.core.modelwrapper import ModelWrapper
-from finn.custom_op.general.im2col import compute_conv_output_dim
+from finn.custom_op.general.im2col import compute_conv_output_dim_2D_padding
 from finn.transformation.infer_datatypes import InferDataTypes
 from finn.transformation.infer_shapes import InferShapes
 
@@ -23,10 +23,12 @@ def check_two_dict_for_equality(dict1, dict2):
 
 
 def execution_im2col(
-    x, idt, k_H, k_W, stride, ifm_ch, ifm_dim_H, ifm_dim_W, pad_amt=0, pad_val=0
+    x, idt, k_H, k_W, stride, ifm_ch, ifm_dim_H, ifm_dim_W, pad_amt, pad_val=0
 ):
-    ofm_dim_H = compute_conv_output_dim(ifm_dim_H, k_H, stride, pad_amt)
-    ofm_dim_W = compute_conv_output_dim(ifm_dim_W, k_W, stride, pad_amt)
+    pad_amt_H = pad_amt[0] + pad_amt[2]
+    pad_amt_W = pad_amt[1] + pad_amt[3]
+    ofm_dim_H = compute_conv_output_dim_2D_padding(ifm_dim_H, k_H, stride, pad_amt_H)
+    ofm_dim_W = compute_conv_output_dim_2D_padding(ifm_dim_W, k_W, stride, pad_amt_W)
 
     # set up onnx model
     inp = helper.make_tensor_value_info(
@@ -103,11 +105,13 @@ def test_im2col():
     ifm_ch = 1
     ifm_dim_H = 4
     ifm_dim_W = 4
-    pad_amt = 0
+    pad_amt = [0, 0, 0, 0]
+    pad_amt_H = pad_amt[0] + pad_amt[2]
+    pad_amt_W = pad_amt[1] + pad_amt[3]
     pad_val = 0
 
-    ofm_dim_H = compute_conv_output_dim(ifm_dim_H, k_H, stride, pad_amt)
-    ofm_dim_W = compute_conv_output_dim(ifm_dim_W, k_W, stride, pad_amt)
+    ofm_dim_H = compute_conv_output_dim_2D_padding(ifm_dim_H, k_H, stride, pad_amt_H)
+    ofm_dim_W = compute_conv_output_dim_2D_padding(ifm_dim_W, k_W, stride, pad_amt_W)
 
     x = np.asarray(
         [
@@ -188,11 +192,13 @@ def test_im2col():
     ifm_ch = 2
     ifm_dim_H = 4
     ifm_dim_W = 4
-    pad_amt = 0
+    pad_amt = [0, 0, 0, 0]
+    pad_amt_H = pad_amt[0] + pad_amt[2]
+    pad_amt_W = pad_amt[1] + pad_amt[3]
     pad_val = 0
 
-    ofm_dim_H = compute_conv_output_dim(ifm_dim_H, k_H, stride, pad_amt)
-    ofm_dim_W = compute_conv_output_dim(ifm_dim_W, k_W, stride, pad_amt)
+    ofm_dim_H = compute_conv_output_dim_2D_padding(ifm_dim_H, k_H, stride, pad_amt_H)
+    ofm_dim_W = compute_conv_output_dim_2D_padding(ifm_dim_W, k_W, stride, pad_amt_W)
 
     x = np.asarray(
         [
@@ -244,11 +250,13 @@ def test_im2col():
     ifm_ch = 2
     ifm_dim_H = 4
     ifm_dim_W = 4
-    pad_amt = 1
+    pad_amt = [1, 1, 1, 1]
+    pad_amt_H = pad_amt[0] + pad_amt[2]
+    pad_amt_W = pad_amt[1] + pad_amt[3]
     pad_val = 0
 
-    ofm_dim_H = compute_conv_output_dim(ifm_dim_H, k_H, stride, pad_amt)
-    ofm_dim_W = compute_conv_output_dim(ifm_dim_W, k_W, stride, pad_amt)
+    ofm_dim_H = compute_conv_output_dim_2D_padding(ifm_dim_H, k_H, stride, pad_amt_H)
+    ofm_dim_W = compute_conv_output_dim_2D_padding(ifm_dim_W, k_W, stride, pad_amt_W)
 
     x = np.asarray(
         [
@@ -320,11 +328,13 @@ def test_im2col():
     ifm_ch = 2
     ifm_dim_H = 4
     ifm_dim_W = 5
-    pad_amt = 0
+    pad_amt = [0, 0, 0, 0]
+    pad_amt_H = pad_amt[0] + pad_amt[2]
+    pad_amt_W = pad_amt[1] + pad_amt[3]
     pad_val = 0
 
-    ofm_dim_H = compute_conv_output_dim(ifm_dim_H, k_H, stride, pad_amt)
-    ofm_dim_W = compute_conv_output_dim(ifm_dim_W, k_W, stride, pad_amt)
+    ofm_dim_H = compute_conv_output_dim_2D_padding(ifm_dim_H, k_H, stride, pad_amt_H)
+    ofm_dim_W = compute_conv_output_dim_2D_padding(ifm_dim_W, k_W, stride, pad_amt_W)
 
     x = np.asarray(
         [
@@ -379,11 +389,13 @@ def test_im2col():
     ifm_ch = 2
     ifm_dim_H = 4
     ifm_dim_W = 5
-    pad_amt = 0
+    pad_amt = [0, 0, 0, 0]
+    pad_amt_H = pad_amt[0] + pad_amt[2]
+    pad_amt_W = pad_amt[1] + pad_amt[3]
     pad_val = 0
 
-    ofm_dim_H = compute_conv_output_dim(ifm_dim_H, k_H, stride, pad_amt)
-    ofm_dim_W = compute_conv_output_dim(ifm_dim_W, k_W, stride, pad_amt)
+    ofm_dim_H = compute_conv_output_dim_2D_padding(ifm_dim_H, k_H, stride, pad_amt_H)
+    ofm_dim_W = compute_conv_output_dim_2D_padding(ifm_dim_W, k_W, stride, pad_amt_W)
 
     x = np.asarray(
         [
@@ -432,11 +444,13 @@ def test_im2col():
     ifm_ch = 2
     ifm_dim_H = 4
     ifm_dim_W = 5
-    pad_amt = 1
+    pad_amt = [1, 1, 1, 1]
+    pad_amt_H = pad_amt[0] + pad_amt[2]
+    pad_amt_W = pad_amt[1] + pad_amt[3]
     pad_val = 0
 
-    ofm_dim_H = compute_conv_output_dim(ifm_dim_H, k_H, stride, pad_amt)
-    ofm_dim_W = compute_conv_output_dim(ifm_dim_W, k_W, stride, pad_amt)
+    ofm_dim_H = compute_conv_output_dim_2D_padding(ifm_dim_H, k_H, stride, pad_amt_H)
+    ofm_dim_W = compute_conv_output_dim_2D_padding(ifm_dim_W, k_W, stride, pad_amt_W)
 
     x = np.asarray(
         [
@@ -505,11 +519,13 @@ def test_im2col():
     ifm_ch = 2
     ifm_dim_H = 5
     ifm_dim_W = 1
-    pad_amt = 0
+    pad_amt = [0, 0, 0, 0]
+    pad_amt_H = pad_amt[0] + pad_amt[2]
+    pad_amt_W = pad_amt[1] + pad_amt[3]
     pad_val = 0
 
-    ofm_dim_H = compute_conv_output_dim(ifm_dim_H, k_H, stride, pad_amt)
-    ofm_dim_W = compute_conv_output_dim(ifm_dim_W, k_W, stride, pad_amt)
+    ofm_dim_H = compute_conv_output_dim_2D_padding(ifm_dim_H, k_H, stride, pad_amt_H)
+    ofm_dim_W = compute_conv_output_dim_2D_padding(ifm_dim_W, k_W, stride, pad_amt_W)
 
     x = np.asarray(
         [[[[1, -1]], [[2, -2]], [[3, -3]], [[4, -4]], [[5, -5]]]],
@@ -536,11 +552,13 @@ def test_im2col():
     ifm_ch = 2
     ifm_dim_H = 5
     ifm_dim_W = 1
-    pad_amt = 1
+    pad_amt = [1, 0, 1, 0]
+    pad_amt_H = pad_amt[0] + pad_amt[2]
+    pad_amt_W = pad_amt[1] + pad_amt[3]
     pad_val = 0
 
-    ofm_dim_H = compute_conv_output_dim(ifm_dim_H, k_H, stride, pad_amt)
-    ofm_dim_W = compute_conv_output_dim(ifm_dim_W, k_W, stride, pad_amt)
+    ofm_dim_H = compute_conv_output_dim_2D_padding(ifm_dim_H, k_H, stride, pad_amt_H)
+    ofm_dim_W = compute_conv_output_dim_2D_padding(ifm_dim_W, k_W, stride, pad_amt_W)
 
     x = np.asarray(
         [[[[1, -1]], [[2, -2]], [[3, -3]], [[4, -4]], [[5, -5]]]],
@@ -575,11 +593,13 @@ def test_im2col():
     ifm_ch = 2
     ifm_dim_H = 5
     ifm_dim_W = 1
-    pad_amt = 1
+    pad_amt = [1, 0, 1, 0]
+    pad_amt_H = pad_amt[0] + pad_amt[2]
+    pad_amt_W = pad_amt[1] + pad_amt[3]
     pad_val = 0
 
-    ofm_dim_H = compute_conv_output_dim(ifm_dim_H, k_H, stride, pad_amt)
-    ofm_dim_W = compute_conv_output_dim(ifm_dim_W, k_W, stride, pad_amt)
+    ofm_dim_H = compute_conv_output_dim_2D_padding(ifm_dim_H, k_H, stride, pad_amt_H)
+    ofm_dim_W = compute_conv_output_dim_2D_padding(ifm_dim_W, k_W, stride, pad_amt_W)
 
     x = np.asarray(
         [[[[1, -1]], [[2, -2]], [[3, -3]], [[4, -4]], [[5, -5]]]],
@@ -607,10 +627,12 @@ def test_im2col_infer_shapes():
     ifm_ch = 1
     ifm_dim_H = 4
     ifm_dim_W = 4
-    pad_amt = 0  # default
+    pad_amt = [0, 0, 0, 0]  # default
+    pad_amt_H = pad_amt[0] + pad_amt[2]
+    pad_amt_W = pad_amt[1] + pad_amt[3]
 
-    ofm_dim_H = compute_conv_output_dim(ifm_dim_H, k_H, stride, pad_amt)
-    ofm_dim_W = compute_conv_output_dim(ifm_dim_W, k_W, stride, pad_amt)
+    ofm_dim_H = compute_conv_output_dim_2D_padding(ifm_dim_H, k_H, stride, pad_amt_H)
+    ofm_dim_W = compute_conv_output_dim_2D_padding(ifm_dim_W, k_W, stride, pad_amt_W)
 
     # set up onnx model
     inp = helper.make_tensor_value_info(