Caffe2 - Python API
A deep learning, cross platform ML framework
pixelshuffle.py
1 from .module import Module
2 from .. import functional as F
3 from ..._jit_internal import weak_module, weak_script_method
4 
5 
6 @weak_module
7 class PixelShuffle(Module):
8  r"""Rearranges elements in a tensor of shape :math:`(*, C \times r^2, H, W)`
9  to a tensor of shape :math:`(*, C, H \times r, W \times r)`.
10 
11  This is useful for implementing efficient sub-pixel convolution
12  with a stride of :math:`1/r`.
13 
14  Look at the paper:
15  `Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network`_
16  by Shi et. al (2016) for more details.
17 
18  Args:
19  upscale_factor (int): factor to increase spatial resolution by
20 
21  Shape:
22  - Input: :math:`(N, L, H_{in}, W_{in})` where :math:`L=C \times \text{upscale\_factor}^2`
23  - Output: :math:`(N, C, H_{out}, W_{out})` where
24  :math:`H_{out} = H_{in} \times \text{upscale\_factor}`
25  and :math:`W_{out} = W_{in} \times \text{upscale\_factor}`
26 
27  Examples::
28 
29  >>> pixel_shuffle = nn.PixelShuffle(3)
30  >>> input = torch.randn(1, 9, 4, 4)
31  >>> output = pixel_shuffle(input)
32  >>> print(output.size())
33  torch.Size([1, 1, 12, 12])
34 
35  .. _Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network:
36  https://arxiv.org/abs/1609.05158
37  """
38  __constants__ = ['upscale_factor']
39 
40  def __init__(self, upscale_factor):
41  super(PixelShuffle, self).__init__()
42  self.upscale_factor = upscale_factor
43 
44  @weak_script_method
45  def forward(self, input):
46  return F.pixel_shuffle(input, self.upscale_factor)
47 
48  def extra_repr(self):
49  return 'upscale_factor={}'.format(self.upscale_factor)