Skip to content

Functionals

einconv.functionals.convNd 

convNd(x: Tensor, weight: Union[Tensor, Parameter], bias: Union[Tensor, Parameter, None] = None, stride: Union[int, Tuple[int, ...]] = 1, padding: Union[int, str, Tuple[int, ...]] = 0, dilation: Union[int, Tuple[int, ...]] = 1, groups: int = 1, simplify: bool = True) -> Tensor

Generalization of torch.nn.functional.conv{1,2,3}d to Nd.

N is determined from the input tensor: It's first axis is the batch dimension, the second axis the channel dimension, and the remaining number of dimensions is interpreted as spatial dimension (with number of spatial dimensions N)

Parameters:

  • x (Tensor) –

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

  • weight (Union[Tensor, Parameter]) –

    Kernel of the convolution. Has shape [out_channels, in_channels / groups, *kernel_size] where kernel_size is an N-tuple of kernel dimensions.

  • bias (Union[Tensor, Parameter, None], default: None ) –

    Optional bias vector of the convolution. Has shape [out_channels]. Default: None.

  • 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, str, 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'.

  • 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.

  • simplify (bool, default: True ) –

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

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/functionals/conv.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def convNd(
    x: Tensor,
    weight: Union[Tensor, Parameter],
    bias: Union[Tensor, Parameter, None] = None,
    stride: Union[int, Tuple[int, ...]] = 1,
    padding: Union[int, str, Tuple[int, ...]] = 0,
    dilation: Union[int, Tuple[int, ...]] = 1,
    groups: int = 1,
    simplify: bool = True,
) -> Tensor:
    """Generalization of ``torch.nn.functional.conv{1,2,3}d`` to ``N``d.

    ``N`` is determined from the input tensor: It's first axis is the batch dimension,
    the second axis the channel dimension, and the remaining number of dimensions is
    interpreted as spatial dimension (with number of spatial dimensions ``N``)

    Args:
        x: Convolution input. Has shape ``[batch_size, in_channels, *input_sizes]``
            where ``len(input_sizes) == N``.
        weight: Kernel of the convolution. Has shape ``[out_channels,
            in_channels / groups, *kernel_size]`` where ``kernel_size`` is an
            ``N``-tuple of kernel dimensions.
        bias: Optional bias vector of the convolution. Has shape ``[out_channels]``.
            Default: ``None``.
        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``.
        simplify: Whether to use a simplified einsum expression. Default: ``True``.

    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 ...``.
    """
    _check_args(x, weight, bias=bias, groups=groups)
    equation, operands, shape = convNd_forward.einsum_expression(
        x,
        weight,
        stride=stride,
        padding=padding,
        dilation=dilation,
        groups=groups,
        simplify=simplify,
    )
    output = einsum(equation, *operands).reshape(shape)

    if bias is not None:
        N = x.dim() - 2
        out_channels = weight.shape[0]
        shape_before_expand = (1, out_channels) + N * (1,)
        output += bias.reshape(*shape_before_expand).expand_as(output)

    return output

einconv.functionals.unfoldNd 

unfoldNd(x: Tensor, kernel_size: Union[int, Tuple[int, ...]], dilation: Union[int, Tuple[int, ...]] = 1, padding: Union[int, Tuple[int, ...], str] = 0, stride: Union[int, Tuple[int, ...]] = 1, simplify: bool = True) -> Tensor

Torch functional for N-dimensional input unfolding that uses einsum.

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

Accepts batched tensors with N spatial dimensions. Acts like torch.nn.functional.unfold for a 4d input (batched images), but works for arbitrary N. See https://pytorch.org/docs/stable/nn.functional.html#unfold.

Parameters:

  • x (Tensor) –

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

  • 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, ...], 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'.

  • 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.

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/functionals/unfold.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def unfoldNd(
    x: Tensor,
    kernel_size: Union[int, Tuple[int, ...]],
    dilation: Union[int, Tuple[int, ...]] = 1,
    padding: Union[int, Tuple[int, ...], str] = 0,
    stride: Union[int, Tuple[int, ...]] = 1,
    simplify: bool = True,
) -> Tensor:
    """Torch functional for N-dimensional input unfolding that uses einsum.

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

    Accepts batched tensors with ``N`` spatial dimensions. Acts like
    ``torch.nn.functional.unfold`` for a 4d input (batched images), but works for
    arbitrary ``N``. See https://pytorch.org/docs/stable/nn.functional.html#unfold.

    Args:
        x: Convolution input. Has shape ``[batch_size, in_channels, *input_sizes]``
            where ``len(input_sizes) == N``.
        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``.

    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 ...)``.
    """
    equation, operands, shape = convNd_unfold.einsum_expression(
        x,
        kernel_size,
        stride=stride,
        padding=padding,
        dilation=dilation,
        simplify=simplify,
    )
    return einsum(equation, *operands).reshape(shape)