class

PixelShuffle

extendsModule
PixelShuffle(upscale_factor: int)
source

Rearrange channel-stacked sub-pixel planes into a high-resolution output.

PixelShuffle implements the sub-pixel convolution (ESPCN) operation: given an input with Cr2C \cdot r^2 channels and spatial size (H,W)(H, W), it reshapes and transposes to produce an output with CC channels and spatial size (Hr,Wr)(H \cdot r, W \cdot r).

output[n,c,Hsh+i,Wsw+j]=input[n,cr2+shr+sw,H,W]\text{output}[n,\, c,\, H \cdot s_h + i,\, W \cdot s_w + j] = \text{input}[n,\, c \cdot r^2 + s_h \cdot r + s_w,\, H,\, W]

where sh,sw{0,,r1}s_h, s_w \in \{0, \dots, r-1\} index the sub-pixel offset and cc indexes the output channel.

Parameters

upscale_factorint
Spatial upscaling factor rr. The number of input channels must be divisible by r2r^2.

Attributes

upscale_factorint
Stored value of the upscale_factor constructor argument.

Notes

  • Input: (N,Cr2,H,W)(N, C \cdot r^2, H, W).
  • Output: (N,C,Hr,Wr)(N, C, H \cdot r, W \cdot r).
  • PixelShuffle is the core operation in the Efficient Sub-Pixel CNN (ESPCN) super-resolution architecture (Shi et al., 2016). The idea is to learn the upsampling entirely through convolution on the low-resolution grid, then rearrange the r2r^2 feature channels into a single high-resolution plane, which is cheaper than computing features at high resolution.
  • Implemented as reshape → permute → reshape on the C++ engine; no intermediate copies beyond those required by the permute.
  • PixelUnshuffle is the exact inverse.

Examples

**2× super-resolution decoder block:**
>>> import lucid
>>> import lucid.nn as nn
>>>
>>> r = 2
>>> # Conv outputs C*r^2 channels at LR spatial size, PixelShuffle expands
>>> block = nn.Sequential(
...     nn.Conv2d(64, 3 * r * r, kernel_size=3, padding=1),
...     nn.PixelShuffle(upscale_factor=r),
... )
>>> x = lucid.zeros(1, 64, 56, 56)
>>> block(x).shape
(1, 3, 112, 112)
**Verify channel / spatial trade-off:**
>>> ps = nn.PixelShuffle(upscale_factor=3)
>>> x = lucid.zeros(2, 36, 10, 10)   # 36 = 4 * 3^2
>>> ps(x).shape
(2, 4, 30, 30)

Methods (3)

dunder

__init__

None
__init__(upscale_factor: int)
source

Initialise the PixelShuffle module. See the class docstring for parameter semantics.

fn

forward

Tensor
forward(x: Tensor)
source

Upsample the input tensor.

Parameters

inputTensor
Input tensor of shape (N,C,)(N, C, *).

Returns

Tensor

Upsampled output tensor.

fn

extra_repr

str
extra_repr()
source

Return a string representation of the layer's configuration.