Pixel Shuffle (Channel to Pixel)

medium (<10 mins) einops super-resolution pixel shuffle

this year by E

Description

The inverse of pixel unshuffle, also known as depth-to-space. It is used to upscale an image by rearranging elements from the channel dimension into spatial blocks. Given a tensor of shape (B, C, H, W) and an upscale factor S, transform it to (B, C // (S*S), H * S, W * S).

Starter Code

import torch
from einops import rearrange

def pixel_shuffle(tensor, upscale_factor):
    # Your einops code here
    pass

Verification

For an input of shape (10, 12, 112, 112) and an upscale_factor of 2, the output should have the shape (10, 3, 224, 224).