ML Katas

Description

Depth-to-space is the inverse of the space-to-depth operation. It rearranges features from the channel dimension into spatial blocks, increasing spatial resolution and decreasing channel depth. For an input of (B, C * S*S, H, W) and a block size S, the output will be (B, C, H * S, W * S).

Guidance

Use einops.rearrange. The pattern needs to decompose the channel dimension into the original channels C and the two block dimensions (s1, s2). Then, move s1 and s2 to become part of the height and width dimensions, respectively.

Starter Code

import torch
from einops import rearrange

def depth_to_space(x, block_size):
    # x: (B, C * S * S, H, W)
    # Your rearrange pattern here
    output = rearrange(x, 'b (c s1 s2) h w -> b c (h s1) (w s2)', s1=block_size, s2=block_size)
    return output

Verification

For an input of shape (10, 12, 112, 112) and block_size=2, the output shape should be (10, 3, 224, 224).