在pytorch中，有什么向量化的方法可以代替两个FOR循环来完成这个操作？

Question

你好，我有一个张量'A'在Pytorch的dimesnsions Batch x Channel x Height x Width。我想把它重塑成'B'，使二元组H和W增加'r'，通道减少'r^2'的系数。在图中，如果'B'有4个通道，那么前4个通道，第4个通道的边框颜色为紫色，那么'A'中的前4个通道将是边框颜色为红、绿、蓝、紫的peackskin色像素，其余通道也会相应调整。

我知道'pack'和'unpack'各可以用2个for循环来完成。但这需要更多的时间。在PyTorch中应该有一种矢量化的方法，只需使用reshape和pemutation命令就可以在 "A "和 "B "之间进行切换。谁能帮我解决这个问题？

在本例中，批处理尺寸设置为 1，但如果批处理尺寸较多，我希望图中所示的操作能对每个批处理条目进行单独操作。

谁能帮我提供一个通用代码，在PyTorch中以向量化的方式在A和B之间进行切换。如果在Batch大小超过1的时候也能用，那就更好了。

请注意，这两个操作不能用已经实现的 nn.PixelShuffle 来完成。

非常感谢。

Answer 1

虽然这可以通过仔细的排列和重塑来实现，但pytorch已经用以下方法实现了这一点 nn.PixelShuffle.