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

Question

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

Answer 1

虽然torch.nn.function.pixel_shuffle 的作用完全一样tf.nn.depth_to_space，但 PyTorch 没有任何函数来执行类似于 的逆操作tf.nn.space_to_depth。

space_to_depth话虽这么说，使用torch.nn.function.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)