output shapes of the more common convolutions and pooling operations in deep learning

out_shape.py

import math


def conv1d_out_shape(input_size, kernel_size, stride=1, padding=0, dilation=1):
    output_size = ((input_size + (2 * padding[0]) - dilation[0] * (kernel_size[0] - 1) - 1) / stride[0]) + 1
    return math.floor(output_size)


def transpose_conv1d_out_shape(input_size, kernel_size, stride=1, padding=0, dilation=1):
    output_size = ((input_size - 1) * stride) - (2 * padding) + (dilation * (kernel_size - 1))
    return output_size


def unpooling1d_out_shape(input_size, kernel_size, stride=None, padding=0):
    if stride is None:
        stride = kernel_size

    output_size = ((input_size - 1) * stride) - (2 * padding) + kernel_size
    return output_size


def conv2d_out_shape(input_size, in_channels, out_channels, kernel_size, stride=(1, 1), padding=(0, 0), dilation=(1, 1), batch_size=None):
    height_in, width_in = input_size
    if type(kernel_size) != list:
        kernel_size = [kernel_size] * 2
    if type(stride) != list:
        stride = [stride] * 2
    if type(padding) != list:
        padding = [padding] * 2
    if type(dilation) != list:
        dilation = [dilation] * 2

    height_out = conv1d_out_shape(height_in, kernel_size, stride, padding, dilation)
    width_out = conv1d_out_shape(width_in, kernel_size, stride, padding, dilation)

    output_shape = (out_channels, height_out, width_out)

    if batch_size and type(batch_size) == int:
        output_shape.insert(0, batch_size)

    return output_shape


def pooling2d_out_shape(input_size, in_channels, out_channels, kernel_size, stride=None, padding=(0, 0), dilation=(1, 1), batch_size=None):
    if stride is None:
        stride = kernel_size

    output_shape = conv2d_out_shape(input_size, in_channels, out_channels, kernel_size, stride, padding, dilation, batch_size)
    return output_shape


def transposed_conv2d_out_shape(input_size, in_channels, out_channels, kernel_size, stride=(1, 1), padding=(0, 0), dilation=(1, 1), batch_size=None):
    height_in, width_in = input_size
    if type(kernel_size) != list:
        kernel_size = [kernel_size] * 2
    if type(stride) != list:
        stride = [stride] * 2
    if type(padding) != list:
        padding = [padding] * 2
    if type(dilation) != list:
        dilation = [dilation] * 2

    height_out = transpose_conv1d_out_shape(height_in, kernel_size, stride, padding, dilation)
    width_out = transpose_conv1d_out_shape(width_in, kernel_size, stride, padding, dilation)

    output_shape = (out_channels, height_out, width_out)

    if batch_size and type(batch_size) == int:
        output_shape.insert(0, batch_size)

    return output_shape


def unpooling2d_out_shape(input_size, in_channels, kernel_size, stride=None, padding=0, batch_size=False):
    height_in, width_in = input_size

    if stride is None:
        stride = kernel_size

    height_out = unpooling1d_out_shape(height_in, kernel_size, stride, padding)
    width_out = unpooling1d_out_shape(width_in, kernel_size, stride, padding)

    output_shape = (in_channels, height_out, width_out)

    if batch_size and type(batch_size) == int:
        output_shape.insert(0, batch_size)

    return output_shape