tf.nn.space_to_depth是否有等效的PyTorch函数？

Question

正如标题所述，tf.nn.space_to_depth是否具有等效的PyTorch函数？

Answer 1

尽管torch.nn.functional.pixel_shuffle确实完成了tf.nn.depth_to_space的工作，但PyTorch没有类似于tf.nn.space_to_depth的任何功能来执行逆运算。

话虽如此，使用torch.nn.functional.unfold来实现space_to_depth很容易。

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)