pytorch上space2depth的实现

申明：转载注明出处
https://www.cnblogs.com/wioponsen/p/14570312.html

方式一 使用torch.nn.functional.unfold

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

方式二 使用view+permute

def space_to_depth(in_tensor, down_scale):
    Batchsize, Ch, Height, Width = in_tensor.size()
    out_channel = Ch * (down_scale ** 2)
    out_Height = Height // down_scale
    out_Width = Width // down_scale

    in_tensor_view = in_tensor.view(Batchsize * Ch, out_Height, down_scale, out_Width, down_scale)
    output = in_tensor_view.permute(0, 2, 4, 1, 3).contiguous().view(Batchsize, out_channel, out_Height, out_Width)
    return output

---

pytorch的pixel_unshuffle 与 tensorflow的Space2depth 的理解与区别

1. 这两个函数都是在做类似的事情：内存重排
2. 两者内存重排在C通道上存在差异：
   a.对于 pytorch [N,C,H,W]，拆分方式为[N,C,H//s,s,W//s,s]，组合为[N,C*s*s, H//s,W//s]
   b.对于 tensorflow [N,H,W,C]，拆分方式为[N,H//s,s,W//s,s,C], 组合为[N,H//s,W//s, s*s*C]
   可以看到，其实都是拆分H和W,然后并到C上，但是，可能是由于tf和pt的排列习惯，在合并新的通道的时候，并没有考虑到这样一个转置的差异：C*s*s 和 s*s*C,所以 pixel_unshuffle 和 Space2depth 的转换只需要在通道上进行一个transpose重排