Function conv_transpose_output_shape

axlearn/common/convolution.py:1077–1138

Returns output size for conv transpose. Each mode follows the formulas below, * SAME: padding=(min(window-1, ceil((w+s-2)/2)), max(stride-1, floor((w+s-2)/2))) pad_total = window+stride-2 output_size = input_size*stride * VALID: padding=(window-1, max(stride-1, window-1)

(
    in_shape: Sequence[Optional[int]],
    *,
    window: Sequence[int],
    strides: Sequence[int],
    padding: ConvPaddingType,
    dilation: Sequence[int],
)

def conv_transpose_output_shape(
    in_shape: Sequence[Optional[int]],
    *,
    window: Sequence[int],
    strides: Sequence[int],
    padding: ConvPaddingType,
    dilation: Sequence[int],
) -> Sequence[int]:
    """Returns output size for conv transpose.

    Each mode follows the formulas below,
    * SAME: padding=(min(window-1, ceil((w+s-2)/2)), max(stride-1, floor((w+s-2)/2)))
        pad_total = window+stride-2
        output_size = input_size*stride
    * VALID: padding=(window-1, max(stride-1, window-1))
        pad_total = window+stride-2 + max(window-stride, 0)
        output_size = input_size*stride + max(window-stride, 0)
    * CAUSAL: padding=(window-1, stride-1)
        pad_total = window+stride-2
        output_size = input_size*stride

    Note: In the above equation, `window` will be replaced with `dilate_window` when dilation > 1.
        dilate_window = (window - 1) * dilation + 1. Check conv_dilate_window()

    Refer to
    https://towardsdatascience.com/understand-transposed-convolutions-and-build-your-own-transposed-convolution-layer-from-scratch-4f5d97b2967

    Args:
        in_shape: convolution lhs shape.
        window: convolution window.
        strides: convolution strides.
        padding: convolution padding.
        dilation: convolution dilation.

    Returns:
        The output shape.

    Raises:
        ValueError: If the length of in_shape, window, strides, and padding are not equal.
    """
    if len(in_shape) != len(window) or len(in_shape) != len(strides):
        raise ValueError(
            f"len(in_shape) = {len(in_shape)} must be equal to "
            f"len(window) = {len(window)} and len(strides) = {len(strides)}"
        )

    window = conv_dilate_window(window=window, dilation=dilation)

    def output_shape(in_shape: Optional[int], window: int, stride: int):
        if in_shape is None:
            return None

        if padding == "SAME":
            return in_shape * stride
        elif padding == "VALID":
            return in_shape * stride + max(window - stride, 0)
        elif padding == "CAUSAL":
            return in_shape * stride