[Feature Request / Proposal] Pixel shuffle layer

[kohr-h]

Upsampling based on pixel shuffling has been proposed in the paper Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network (2016). The complete method combines a convolution layer, the pixel shuffle operation and a specific initialization to get rid of block artifacts. For the initialization details, see Checkerboard artifact free sub-pixel convolution: A note on sub-pixel convolution, resize convolution and convolution resize (2017).

Pixel shuffling in 2D means to reshape a tensor of shape (N, f1*f2*C, H, W) to (N, C, f1*H, f2*W), thereby effectively upscaling the images by (f1, f2).
In MXNet, pixel shuffling could be implemented like this in Python:

def hybrid_forward(self, F, x):
        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

Would this be an interesting addition?

[lanking520]

Hi @kohr-h please feel free to raise a PR to contribute to this operator/example in Python.
@szha @zhreshold WDYT?

[zhreshold]

If it can be achieved simply using combination of reshapes, we can put it in mxnet.gluon.nn.contrib module

[kohr-h]

True, it's just reshaping in the end, although not entirely trivial to make work as hybrid block. Regarding placement, gluon.nn.contrib looks okay to me, given that the core collection of layers seems to be rather slim.

[Mut1nyJD]

Nice I was just looking for this

[Mut1nyJD]

Hmm but somehow I'll get strange result I replaced all my Conv2DTranspose layers with a Conv2D(channels= 2 * 2 * num_channel(*)) and then PixelShuffle2D((2,2))
I took the code from your PR but I'll get pretty horrible results.

(*) num_channels = equal the number of channels that used to be in Conv2DTranspose

[kohr-h]

Possible, I have to do more testing myself. I think you need to swap axes before reshaping.

[Mut1nyJD]

@kohr-h

Btw you might want to have look at the example/gluon/Superpixel it looks to me there someone wrote already a PixelShuffle OP in there. I have not tried it yet to see if it works correctly though.

[kohr-h]

@Mut1nyJD I've updated the code in the PR and cross-checked it against the PyTorch implementation.

[kohr-h]

Btw you might want to have look at the example/gluon/Superpixel it looks to me there someone wrote already a PixelShuffle OP in there. I have not tried it yet to see if it works correctly though.

Thanks for the pointer. Indeed, that's the same operation. It would make sense to drop in the new layer there.