Modules

einconv.modules.ConvNd

ConvNd(N: int, in_channels: int, out_channels: int, kernel_size: Union[int, Tuple[int, ...]], stride: Union[int, Tuple[int, ...]] = 1, padding: Union[int, Tuple[int, ...], str] = 0, dilation: Union[int, Tuple[int, ...]] = 1, groups: int = 1, bias: bool = True, padding_mode: str = 'zeros', device: Union[None, torch.device] = None, dtype: Union[None, torch.dtype] = None, simplify: bool = True)

PyTorch module for N-dimensional convolution that uses einsum.

Initialize N-dimensional convolution layer.

Parameters are initialized using the same convention as PyTorch's convolutions.

The weight has shape [out_channels, in_channels // groups, *kernel_size] with len(kernel_size) == N. The bias has shape [out_channels].

Parameters:

• N (int) –

  Convolution dimension. For N=1,2,3 the layer behaves like PyTorch's nn.Conv{N=1,2,3}d. However, it also works for N>3.

• in_channels (int) –

  Number of input channels.

• out_channels (int) –

  Number of output channels.

• kernel_size (Union[int, Tuple[int, ...]]) –

  Kernel dimensions. Can be a single integer (shared along all spatial dimensions), or an N-tuple of integers.

• stride (Union[int, Tuple[int, ...]], default: 1 ) –

  Stride of the convolution. Can be a single integer (shared along all spatial dimensions), or an N-tuple of integers. Default: 1.

• padding (Union[int, Tuple[int, ...], str], default: 0 ) –

  Padding of the convolution. Can be a single integer (shared along all spatial dimensions), an N-tuple of integers, or a string. Default: 0. Allowed strings are 'same' and 'valid'.

• dilation (Union[int, Tuple[int, ...]], default: 1 ) –

  Dilation of the convolution. Can be a single integer (shared along all spatial dimensions), or an N-tuple of integers. Default: 1.

• groups (int, default: 1 ) –

  In how many groups to split the input channels. Default: 1.

• bias (bool, default: True ) –

  Whether to use a bias. Default: True.

• padding_mode (str, default: 'zeros' ) –

  How to perform padding. Default: 'zeros'. No other modes are supported at the moment.

• device (Union[None, device], default: None ) –

  Device on which the module is initialized.

• dtype (Union[None, dtype], default: None ) –

  Data type assumed by the module.

• simplify (bool, default: True ) –

  Whether to use a simplified einsum expression. Default: True.

Raises:

• NotImplementedError –

  For unsupported padding modes.

• ValueError –

  For invalid combinations of in_channels, out_channels, and groups.

Source code in einconv/modules/conv.py

def __init__(
    self,
    N: int,
    in_channels: int,
    out_channels: int,
    kernel_size: Union[int, Tuple[int, ...]],
    stride: Union[int, Tuple[int, ...]] = 1,
    padding: Union[int, Tuple[int, ...], str] = 0,
    dilation: Union[int, Tuple[int, ...]] = 1,
    groups: int = 1,
    bias: bool = True,
    padding_mode: str = "zeros",
    device: Union[None, torch.device] = None,
    dtype: Union[None, torch.dtype] = None,
    simplify: bool = True,
) -> None:
    """Initialize N-dimensional convolution layer.

    Parameters are initialized using the same convention as PyTorch's convolutions.

    The weight has shape ``[out_channels, in_channels // groups, *kernel_size]``
    with ``len(kernel_size) == N``. The bias has shape ``[out_channels]``.

    Args:
        N: Convolution dimension. For ``N=1,2,3`` the layer behaves like PyTorch's
            ``nn.Conv{N=1,2,3}d``. However, it also works for ``N>3``.
        in_channels: Number of input channels.
        out_channels: Number of output channels.
        kernel_size: Kernel dimensions. Can be a single integer (shared along all
            spatial dimensions), or an ``N``-tuple of integers.
        stride: Stride of the convolution. Can be a single integer (shared along all
            spatial dimensions), or an ``N``-tuple of integers. Default: ``1``.
        padding: Padding of the convolution. Can be a single integer (shared along
            all spatial dimensions), an ``N``-tuple of integers, or a string.
            Default: ``0``. Allowed strings are ``'same'`` and ``'valid'``.
        dilation: Dilation of the convolution. Can be a single integer (shared along
            all spatial dimensions), or an ``N``-tuple of integers. Default: ``1``.
        groups: In how many groups to split the input channels. Default: ``1``.
        bias: Whether to use a bias. Default: ``True``.
        padding_mode: How to perform padding. Default: ``'zeros'``. No other modes
            are supported at the moment.
        device: Device on which the module is initialized.
        dtype: Data type assumed by the module.
        simplify: Whether to use a simplified einsum expression. Default: ``True``.

    Raises:
        NotImplementedError: For unsupported padding modes.
        ValueError: For invalid combinations of ``in_channels``, ``out_channels``,
            and ``groups``.
    """
    super().__init__()

    # basic argument checks
    if padding_mode != "zeros":
        raise NotImplementedError("Only padding_mode='zeros' supported.")

    if groups <= 0:
        raise ValueError("groups must be a positive integer")
    if in_channels % groups != 0:
        raise ValueError(
            f"groups ({groups}) must divide in_channels ({in_channels})."
        )
    if out_channels % groups != 0:
        raise ValueError(
            f"groups ({groups}) must divide out_channels ({out_channels})."
        )

    self.N = N
    self.in_channels = in_channels
    self.out_channels = out_channels
    self.simplify = simplify

    # standardize into N-tuple form if possible
    self.kernel_size = _tuple(kernel_size, N)
    self.stride = _tuple(stride, N)
    self.padding = padding if isinstance(padding, str) else _tuple(padding, N)
    self.padding_mode = padding_mode
    self.dilation = _tuple(dilation, N)
    self.groups = groups

    # initialize parameters
    device = torch.device("cpu") if device is None else device
    dtype = torch.float32 if dtype is None else dtype
    weight_shape = (out_channels, in_channels // groups, *self.kernel_size)
    self.register_parameter(
        "weight", Parameter(empty(weight_shape, device=device, dtype=dtype))
    )
    bias_shape = (out_channels,)
    self.register_parameter(
        "bias",
        Parameter(empty(bias_shape, device=device, dtype=dtype)) if bias else None,
    )

    self.reset_parameters()

forward

forward(x: Tensor) -> Tensor

Perform convolution on the input.

Parameters:

• x (Tensor) –

  Convolution input. Has shape [batch_size, in_channels, *input_sizes] where len(input_sizes) == N.

Returns:

• Tensor –

  Result of the convolution. Has shape [batch_size, out_channels, *output_sizes] where len(output_sizes) == N. In einops notation, the index structure is n (g c_out) o1 o2 ....

Source code in einconv/modules/conv.py

def forward(self, x: Tensor) -> Tensor:
    """Perform convolution on the input.

    Args:
        x: Convolution input. Has shape ``[batch_size, in_channels, *input_sizes]``
            where ``len(input_sizes) == N``.

    Returns:
        Result of the convolution. Has shape \
        ``[batch_size, out_channels, *output_sizes]`` where \
        ``len(output_sizes) == N``. In ``einops`` notation, the index structure \
        is ``n (g c_out) o1 o2 ...``.
    """
    return convNd(
        x,
        self.weight,
        bias=self.bias,
        stride=self.stride,
        padding=self.padding,
        dilation=self.dilation,
        groups=self.groups,
        simplify=self.simplify,
    )

from_nn_Conv classmethod

from_nn_Conv(conv_module: Union[Conv1d, Conv2d, Conv3d], simplify: bool = True) -> ConvNd

Convert a torch.nn.Conv{1,2,3}d module to a ConvNd layer.

Parameters:

• conv_module (Union[Conv1d, Conv2d, Conv3d]) –

  Convolution module.

• simplify (bool, default: True ) –

  Whether to use a simplified einsum expression. Default: True.

Returns:

• ConvNd –

  Converted ConvNd module.

Source code in einconv/modules/conv.py

@classmethod
def from_nn_Conv(
    cls, conv_module: Union[Conv1d, Conv2d, Conv3d], simplify: bool = True
) -> ConvNd:
    """Convert a ``torch.nn.Conv{1,2,3}d`` module to a ``ConvNd`` layer.

    Args:
        conv_module: Convolution module.
        simplify: Whether to use a simplified einsum expression. Default: ``True``.

    Returns:
        Converted ConvNd module.
    """
    N = {Conv1d: 1, Conv2d: 2, Conv3d: 3}[conv_module.__class__]
    einconv_module = cls(
        N,
        conv_module.in_channels,
        conv_module.out_channels,
        conv_module.kernel_size,
        stride=conv_module.stride,
        padding=conv_module.padding,
        dilation=conv_module.dilation,
        groups=conv_module.groups,
        bias=conv_module.bias is not None,
        padding_mode=conv_module.padding_mode,
        device=conv_module.weight.device,
        dtype=conv_module.weight.dtype,
        simplify=simplify,
    )
    sync_parameters(conv_module, einconv_module)

    return einconv_module

einconv.modules.UnfoldNd

UnfoldNd(kernel_size: Union[int, Tuple[int, ...]], dilation: Union[int, Tuple[int, ...]] = 1, padding: Union[int, Tuple[int, ...]] = 0, stride: Union[int, Tuple[int, ...]] = 1, simplify: bool = True)

PyTorch module for N-dimensional input unfolding (im2col) that uses einsum.

Extracts sliding local blocks from a batched input tensor (im2col).

This module accepts batched tensors with N spatial dimensions. It acts like torch.nn.Unfold for a 4d input (batched images), but works for arbitrary N. See https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html.

Parameters:

• kernel_size (Union[int, Tuple[int, ...]]) –

  Kernel dimensions. Can be a single integer (shared along all spatial dimensions), or an N-tuple of integers.

• dilation (Union[int, Tuple[int, ...]], default: 1 ) –

  Dilation of the convolution. Can be a single integer (shared along all spatial dimensions), or an N-tuple of integers. Default: 1.

• padding (Union[int, Tuple[int, ...]], default: 0 ) –

  Padding of the convolution. Can be a single integer (shared along all spatial dimensions), an N-tuple of integers, or a string. Default: 0. Allowed strings are 'same' and 'valid'.

• stride (Union[int, Tuple[int, ...]], default: 1 ) –

  Stride of the convolution. Can be a single integer (shared along all spatial dimensions), or an N-tuple of integers. Default: 1.

• simplify (bool, default: True ) –

  Whether to use a simplified einsum expression. Default: True.

Source code in einconv/modules/unfold.py

def __init__(
    self,
    kernel_size: Union[int, Tuple[int, ...]],
    dilation: Union[int, Tuple[int, ...]] = 1,
    padding: Union[int, Tuple[int, ...]] = 0,
    stride: Union[int, Tuple[int, ...]] = 1,
    simplify: bool = True,
):
    """Extracts sliding local blocks from a batched input tensor (``im2col``).

    This module accepts batched tensors with ``N`` spatial dimensions. It acts like
    ``torch.nn.Unfold`` for a 4d input (batched images), but works for arbitrary
    ``N``. See https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html.

    Args:
        kernel_size: Kernel dimensions. Can be a single integer (shared along all
            spatial dimensions), or an ``N``-tuple of integers.
        dilation: Dilation of the convolution. Can be a single integer (shared along
            all spatial dimensions), or an ``N``-tuple of integers. Default: ``1``.
        padding: Padding of the convolution. Can be a single integer (shared along
            all spatial dimensions), an ``N``-tuple of integers, or a string.
            Default: ``0``. Allowed strings are ``'same'`` and ``'valid'``.
        stride: Stride of the convolution. Can be a single integer (shared along all
            spatial dimensions), or an ``N``-tuple of integers. Default: ``1``.
        simplify: Whether to use a simplified einsum expression. Default: ``True``.
    """
    super().__init__()

    self._kernel_size = kernel_size
    self._dilation = dilation
    self._padding = padding
    self._stride = stride
    self._simplify = simplify

forward

forward(x: Tensor) -> Tensor

Compute the unfolded input.

Parameters:

• x (Tensor) –

  Convolution input. Has shape [batch_size, in_channels, *input_sizes] where len(input_sizes) == N.

Returns:

• Tensor –

  Unfolded input. Has shape [batch_size, in_channels * tot_kernel_size, tot_output_size] where tot_kernel_size is the kernel dimension product and tot_output_size is the product of the output spatial dimensions. In einops notation, the index structure is n (c_in k1 k2 ...) (o1 o2 ...).

Source code in einconv/modules/unfold.py

def forward(self, x: Tensor) -> Tensor:
    """Compute the unfolded input.

    Args:
        x: Convolution input. Has shape ``[batch_size, in_channels, *input_sizes]``
            where ``len(input_sizes) == N``.

    Returns:
        Unfolded input. Has shape \
        ``[batch_size, in_channels * tot_kernel_size, tot_output_size]`` where \
        ``tot_kernel_size`` is the kernel dimension product and \
        ``tot_output_size`` is the product of the output spatial dimensions. In \
        ``einops`` notation, the index structure is \
        ``n (c_in k1 k2 ...) (o1 o2 ...)``.
    """
    return unfoldNd(
        x,
        self._kernel_size,
        dilation=self._dilation,
        padding=self._padding,
        stride=self._stride,
        simplify=self._simplify,
    )