torch.nn.modules.fold 的源代码

from .module import Module
from .. import functional as F

from torch import Tensor
from ..common_types import _size_any_t

__all__ = ['Fold', 'Unfold']

class Fold(Module):
    r"""将一个滑动局部块数组组合成一个大的包含张量。

    考虑一个包含滑动局部块的批量 :attr:`input` 张量，
    例如，图像的补丁，形状为 :math:`(N, C \times  \prod(\text{kernel\_size}), L)`，
    其中 :math:`N` 是批次维度，:math:`C \times \prod(\text{kernel\_size})`
    是块内的值的数量（一个块有 :math:`\prod(\text{kernel\_size})`
    空间位置，每个位置包含一个 :math:`C` 通道的向量），并且
    :math:`L` 是块的总数。（这与 :class:`~torch.nn.Unfold` 的输出形状完全相同。）此
    操作将这些局部块组合成大的 :attr:`output` 张量
    形状为 :math:`(N, C, \text{output\_size}[0], \text{output\_size}[1], \dots)`
    通过求和重叠的值。与 :class:`~torch.nn.Unfold` 类似，参数必须满足

    .. math::
        L = \prod_d \left\lfloor\frac{\text{output\_size}[d] + 2 \times \text{padding}[d] %
            - \text{dilation}[d] \times (\text{kernel\_size}[d] - 1) - 1}{\text{stride}[d]} + 1\right\rfloor,

    其中 :math:`d` 遍历所有空间维度。

    * :attr:`output_size` 描述滑动局部块的大包含张量的空间形状。
      当多个输入形状映射到相同数量的滑动块时，它有助于消除歧义，例如，
      当 ``stride > 0`` 时。

    :attr:`padding`、:attr:`stride` 和 :attr:`dilation` 参数指定
    如何检索滑动块。

    * :attr:`stride` 控制滑动块的步幅。

    * :attr:`padding` 控制重塑前每个维度两侧的隐式零填充量。

    * :attr:`dilation` 控制核点之间的间距；也称为 à trous 算法。
      这很难描述，但这个 `link`_ 有一个很好的可视化效果，展示了 :attr:`dilation` 的作用。

    参数:
        output_size (int or tuple): 输出空间维度的形状（即 ``output.sizes()[2:]``）
        kernel_size (int or tuple): 滑动块的大小
        dilation (int or tuple, optional): 控制邻域内元素步幅的参数。默认值: 1
        padding (int or tuple, optional): 输入两侧要添加的隐式零填充。默认值: 0
        stride (int or tuple): 输入空间维度中滑动块的步幅。默认值: 1

    * 如果 :attr:`output_size`、:attr:`kernel_size`、:attr:`dilation`、
      :attr:`padding` 或 :attr:`stride` 是 int 或长度为 1 的元组，则
      它们的值将在所有空间维度上复制。

    * 对于两个输出空间维度的情况，此操作有时称为 ``col2im``。

    .. note::
        :class:`~torch.nn.Fold` 通过求和所有包含块的值来计算结果
        大张量中的每个组合值。
        :class:`~torch.nn.Unfold` 通过从大张量复制来提取局部块中的值。
        因此，如果块重叠，它们不是彼此的逆。

        通常，折叠和展开操作的关系如下。考虑使用相同
        参数创建的 :class:`~torch.nn.Fold` 和
        :class:`~torch.nn.Unfold` 实例：

        >>> fold_params = dict(kernel_size=..., dilation=..., padding=..., stride=...)
        >>> fold = nn.Fold(output_size=..., **fold_params)
        >>> unfold = nn.Unfold(**fold_params)

        然后对于任何（支持的）``input`` 张量，以下
        等式成立：

        ::

            fold(unfold(input)) == divisor * input

        其中 ``divisor`` 是一个仅依赖于 ``input`` 的形状
        和数据类型的张量：

        >>> # xdoctest: +SKIP
        >>> input_ones = torch.ones(input.shape, dtype=input.dtype)
        >>> divisor = fold(unfold(input_ones))

        当 ``divisor`` 张量不包含零元素时，则
        ``fold`` 和 ``unfold`` 操作是彼此的逆（直到常数除数）。

    .. warning::
        目前仅支持未批处理的（3D）或批处理的（4D）图像类输出张量。

    形状:
        - 输入: :math:`(N, C \times \prod(\text{kernel\_size}), L)` 或 :math:`(C \times \prod(\text{kernel\_size}), L)`
        - 输出: :math:`(N, C, \text{output\_size}[0], \text{output\_size}[1], \dots)`
          或 :math:`(C, \text{output\_size}[0], \text{output\_size}[1], \dots)` 如上所述

    示例::

        >>> fold = nn.Fold(output_size=(4, 5), kernel_size=(2, 2))
        >>> input = torch.randn(1, 3 * 2 * 2, 12)
        >>> output = fold(input)
        >>> output.size()
        torch.Size([1, 3, 4, 5])

    .. _link:
        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md

    """

    __constants__ = ['output_size', 'kernel_size', 'dilation', 'padding',
                     'stride']
    output_size: _size_any_t
    kernel_size: _size_any_t
    dilation: _size_any_t
    padding: _size_any_t
    stride: _size_any_t

    def __init__(
        self,
        output_size: _size_any_t,
        kernel_size: _size_any_t,
        dilation: _size_any_t = 1,
        padding: _size_any_t = 0,
        stride: _size_any_t = 1
    ) -> None:
        super().__init__()
        self.output_size = output_size
        self.kernel_size = kernel_size
        self.dilation = dilation
        self.padding = padding
        self.stride = stride

    def forward(self, input: Tensor) -> Tensor:
        return F.fold(input, self.output_size, self.kernel_size, self.dilation,
                      self.padding, self.stride)

    def extra_repr(self) -> str:
        return 'output_size={output_size}, kernel_size={kernel_size}, ' \
            'dilation={dilation}, padding={padding}, stride={stride}'.format(
                **self.__dict__
            )


class Unfold(Module</span