torch.nn.modules.pixelshuffle 的源代码

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

from torch import Tensor

__all__ = ['PixelShuffle', 'PixelUnshuffle']

class PixelShuffle(Module):
    r"""根据上采样因子重新排列张量中的元素。

    将形状为 :math:`(*, C \times r^2, H, W)` 的张量重新排列为形状为 :math:`(*, C, H \times r, W \times r)` 的张量，其中 r 是上采样因子。

    这对于以步幅为 :math:`1/r` 实现高效的子像素卷积非常有用。

    请参阅论文：
    `使用高效子像素卷积神经网络实现实时单图像和视频超分辨率`_
    由 Shi 等人 (2016) 了解更多详情。

    参数:
        upscale_factor (int): 增加空间分辨率的因子

    形状:
        - 输入: :math:`(*, C_{in}, H_{in}, W_{in})`，其中 * 是零个或多个批次维度
        - 输出: :math:`(*, C_{out}, H_{out}, W_{out})`，其中

    .. math::
        C_{out} = C_{in} \div \text{upscale\_factor}^2

    .. math::
        H_{out} = H_{in} \times \text{upscale\_factor}

    .. math::
        W_{out} = W_{in} \times \text{upscale\_factor}

    示例::

        >>> pixel_shuffle = nn.PixelShuffle(3)
        >>> input = torch.randn(1, 9, 4, 4)
        >>> output = pixel_shuffle(input)
        >>> print(output.size())
        torch.Size([1, 1, 12, 12])

    .. _使用高效子像素卷积神经网络实现实时单图像和视频超分辨率:
        https://arxiv.org/abs/1609.05158
    """

    __constants__ = ['upscale_factor']
    upscale_factor: int

    def __init__(self, upscale_factor: int) -> None:
        super().__init__()
        self.upscale_factor = upscale_factor

    def forward(self, input: Tensor) -> Tensor:
        return F.pixel_shuffle(input, self.upscale_factor)

    def extra_repr(self) -> str:
        return f'upscale_factor={self.upscale_factor}'


class PixelUnshuffle(Module):
    r"""反转 PixelShuffle 操作。

    通过重新排列元素，将形状为 :math:`(*, C, H \times r, W \times r)` 的张量反转为形状为 :math:`(*, C \times r^2, H, W)` 的张量，其中 r 是下采样因子。

    请参阅论文：
    `使用高效子像素卷积神经网络实现实时单图像和视频超分辨率`_
    由 Shi 等人 (2016) 了解更多详情。

    参数:
        downscale_factor (int): 减少空间分辨率的因子

    形状:
        - 输入: :math:`(*, C_{in}, H_{in}, W_{in})`，其中 * 是零个或多个批次维度
        - 输出: :math:`(*, C_{out}, H_{out}, W_{out})`，其中

    .. math::
        C_{out} = C_{in} \times \text{downscale\_factor}^2

    .. math::
        H_{out} = H_{in} \div \text{downscale\_factor}

    .. math::
        W_{out} = W_{in} \div \text{downscale\_factor}

    示例::

        >>> pixel_unshuffle = nn.PixelUnshuffle(3)
        >>> input = torch.randn(1, 1, 12, 12)
        >>> output = pixel_unshuffle(input)
        >>> print(output.size())
        torch.Size([1, 9, 4, 4])

    .. _使用高效子像素卷积神经网络实现实时单图像和视频超分辨率:
        https://arxiv.org/abs/1609.05158
    """

    __constants__ = ['downscale_factor']
    downscale_factor: int

    def __init__(self, downscale_factor: int) -> None:
        super().__init__()
        self.downscale_factor = downscale_factor

    def forward(self, input: Tensor) -> Tensor:
        return F.pixel_unshuffle(input, self.downscale_factor)

    def extra_repr(self) -> str:
        return f'downscale_factor={self.downscale_factor}'