[FIXED] Is there an equivalent PyTorch function for `tf.nn.space_to_depth`

Issue

As the title says, is there an equivalent PyTorch function for tf.nn.space_to_depth?

Solution

While torch.nn.functional.pixel_shuffle does exactly what tf.nn.depth_to_space does, PyTorch doesn't have any function to do the inverse operation similar to tf.nn.space_to_depth.

That being said, it is easy to implement space_to_depth using torch.nn.functional.unfold.

def space_to_depth(x, block_size):
    n, c, h, w = x.size()
    unfolded_x = torch.nn.functional.unfold(x, block_size, stride=block_size)
    return unfolded_x.view(n, c * block_size ** 2, h // block_size, w // block_size)

Answered By - Priyatham