ThomasDelteil commented on a change in pull request #13571: Add pixelshuffle
layers
URL: https://github.com/apache/incubator-mxnet/pull/13571#discussion_r239994774
##########
File path: python/mxnet/gluon/contrib/nn/basic_layers.py
##########
@@ -235,3 +235,135 @@ def _get_num_devices(self):
def hybrid_forward(self, F, x, gamma, beta, running_mean, running_var):
return F.contrib.SyncBatchNorm(x, gamma, beta, running_mean,
running_var,
name='fwd', **self._kwargs)
+
+class PixelShuffle1D(HybridBlock):
+
+ """Pixel-shuffle layer for upsampling in 1 dimension."""
+
+ def __init__(self, factor):
+ """
+
+ Parameters
+ ----------
+ factor : int or 1-tuple of int
+ Upsampling factor, applied to the ``W`` dimension.
+
+ Inputs:
+ - **data**: Tensor of shape ``(N, f*C, W)``.
+ Outputs:
+ - **out**: Tensor of shape ``(N, C, f*W)``.
+
+ Examples
+ --------
+ >>> pxshuf = PixelShuffle1D(2)
+ >>> x = mx.nd.zeros((1, 8, 3))
+ >>> pxshuf(x).shape
+ (1, 4, 6)
+ """
+ super().__init__()
+ self._factor = int(factor)
+
+ def hybrid_forward(self, F, x):
+ """Perform pixel-shuffling on the input."""
+ f = self._factor
+ # (N, f*C, W)
+ x = F.reshape(x, (0, -4, -1, f, 0)) # (N, C, f, W)
+ x = F.reshape(x, (0, 0, -3)) # (N, C, f*W)
+ return x
+
+ def __repr__(self):
+ return "{}({})".format(self.__class__.__name__, self._factor)
+
+
+class PixelShuffle2D(HybridBlock):
+
+ """Pixel-shuffle layer for upsampling in 2 dimensions."""
+
+ def __init__(self, factor):
+ """
+
+ Parameters
+ ----------
+ factor : int or 2-tuple of int
+ Upsampling factors, applied to the ``H`` and ``W`` dimensions,
+ in that order.
+
+ Inputs:
+ - **data**: Tensor of shape ``(N, f1*f2*C, H, W)``.
+ Outputs:
+ - **out**: Tensor of shape ``(N, C, f1*H, f2*W)``.
+
+ Examples
+ --------
+ >>> pxshuf = PixelShuffle2D((2, 3))
+ >>> x = mx.nd.zeros((1, 12, 3, 5))
+ >>> pxshuf(x).shape
+ (1, 2, 6, 15)
+ """
+ super().__init__()
+ try:
+ self._factors = (int(factor),) * 2
+ except TypeError:
+ self._factors = tuple(int(fac) for fac in factor)
+ assert len(self._factors) == 2, "wrong length
{}".format(len(self._factors))
+
+ def hybrid_forward(self, F, x):
+ """Perform pixel-shuffling on the input."""
+ f1, f2 = self._factors
+ # (N, f1*f2*C, H, W)
+ x = F.reshape(x, (0, -4, -1, f1 * f2, 0, 0)) # (N, C, f1*f2, H, W)
+ x = F.reshape(x, (0, 0, -3, 0)) # (N, C, f1*f2*H, W)
+ x = F.reshape(x, (0, 0, -4, -1, f2, 0)) # (N, C, f1*H, f2, W)
+ x = F.reshape(x, (0, 0, 0, -3)) # (N, C, f1*H, f2*W)
+ return x
+
+ def __repr__(self):
+ return "{}({})".format(self.__class__.__name__, self._factors)
+
+
+class PixelShuffle3D(HybridBlock):
+
+ """Pixel-shuffle layer for upsampling in 3 dimensions."""
+
+ def __init__(self, factor):
+ """
+
+ Parameters
+ ----------
+ factor : int or 3-tuple of int
+ Upsampling factors, applied to the ``D``, ``H`` and ``W``
+ dimensions, in that order.
+
+ Inputs:
+ - **data**: Tensor of shape ``(N, f1*f2*f3*C, D, H, W)``.
+ Outputs:
+ - **out**: Tensor of shape ``(N, C, f1*D,f2*H, f3*W)``.
+
+ Examples
+ --------
+ >>> pxshuf = PixelShuffle3D((2, 3, 4))
+ >>> x = mx.nd.zeros((1, 48, 3, 5, 7))
+ >>> pxshuf(x).shape
+ (1, 2, 6, 15, 28)
+ """
+ super().__init__()
+ try:
+ self._factors = (int(factor),) * 3
+ except TypeError:
+ self._factors = tuple(int(fac) for fac in factor)
+ assert len(self._factors) == 3, "wrong length
{}".format(len(self._factors))
+
+ def hybrid_forward(self, F, x):
+ """Perform pixel-shuffling on the input."""
+ f1, f2, f3 = self._factors
+ # (N,
f1*f2*f3*C, D, H, W)
+ x = F.reshape(x, (0, -4, -1, f1 * f2 * f3, 0, 0, 0)) # (N, C,
f1*f2*f3, D, H, W)
Review comment:
Adding tests will let you test that automatically on all our supported
platform when your PR is being tested in the CI. GPU / CPU / Windows / Linux /
arm / etc.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
