diff --git a/src/finn/custom_op/im2col.py b/src/finn/custom_op/im2col.py index 2bd039c3acbcb3dfa99b478b012a4e89d5399f45..16446c15d46ee7996162f864708f7fde6cfedaf3 100644 --- a/src/finn/custom_op/im2col.py +++ b/src/finn/custom_op/im2col.py @@ -11,13 +11,14 @@ from finn.core.datatype import DataType def compute_conv_output_dim(ifm_dim, k, stride, pad=0): - "Return spatial output dimension size for convolution with given params." + """Returns spatial output dimension size for convolution with given params.""" return int(((ifm_dim + 2 * pad - k) / stride) + 1) def get_im2col_indices_nchw( x_shape, field_height, field_width, padding=0, stride_y=1, stride_x=1 ): + """Returns im2col indices.""" # First figure out what the size of the output should be N, C, H, W = x_shape assert (H + 2 * padding - field_height) % stride_y == 0 @@ -41,7 +42,9 @@ def get_im2col_indices_nchw( def im2col_indices_nchw( x, field_height, field_width, padding=0, stride_y=1, stride_x=1, pad_val=0 ): - """ An implementation of im2col based on some fancy indexing """ + """Performs im2col on x with given field height and width, as well as values + for padding and stride size. + Returns result of im2col.""" # Zero-pad the input p = padding x_padded = np.pad(