Pixel Unshuffle (Pixel to Channel)

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

this year by E

Description

This is another name for the space-to-depth operation, common in super-resolution models. It involves rearranging blocks of spatial data into the channel dimension. Given a tensor of shape (B, C, H, W) and a downscale factor S, you need to transform it to (B, C * S * S, H // S, W // S).

Starter Code

import torch
from einops import rearrange

def pixel_unshuffle(tensor, downscale_factor):
    # Your einops code here
    pass

Verification

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