【代码实现】torch实现F.pixel_shuffle和F.pixel_unshuffle
原理
pixel_shuffle 和 pixel_unshuffle 常用于神经网络减少特征图尺寸以减少计算量，由于有些硬件不支持这两个算子，可以根据原理使用torch实现。
代码实现
import torch.nn.functional as F
import torch
def pixelshuffle_inv(tensor, scale=2):N, ch, height, width = tensor.shapenew_ch = ch * (scale * scale)new_height = height // scalenew_width = width // scaletensor = tensor.view(N, ch, new_height, scale, new_width, scale)tensor = tensor.permute(0, 1, 3, 5, 2, 4).contiguous()tensor = tensor.view(N, ch * (scale ** 2), new_height, new_width)return tensordef pixelshuffle(tensor, scale=2):N, ch, height, width = tensor.shapenew_ch = ch // (scale * scale)new_height = height * scalenew_width = width * scale# 重新排列张量output_tensor = tensor.view(N, new_ch, scale, scale, height, width)output_tensor = output_tensor.permute(0, 1, 4, 2, 5, 3).contiguous()output_tensor = output_tensor.view(N, new_ch, new_height, new_width)return output_tensorif __name__ == '__main__':input = torch.randn(1, 3, 256, 256)scale = 2unshuffle_ = pixelshuffle_inv(input,scale)unshuffle_F = F.pixel_unshuffle(input,scale)print(torch.equal(unshuffle_,unshuffle_F))print(torch.max(unshuffle_-unshuffle_F))shuffle_ = pixelshuffle(unshuffle_F,scale)shuffle_F = F.pixel_shuffle(unshuffle_F,scale)print(torch.equal(shuffle_,shuffle_F))print(torch.max(shuffle_-shuffle_F))
运行结果，与官方结果完全一致。
