将pytorch的批处理、高度、宽度、alpha更改为批处理、alpha、高度和宽度

Question

我有一批像这样的图像

torch.Size([10, 512, 512, 3])

我可以循环到图像，可以看到10个图像。但要把这批货送到pytorch，我必须把它转换成

torch.Size([10, 3, 512, 512])

我尝试了很多方法，但无法获得这个的解决方案

我们该怎么做？

Answer 1

使用permute:

import torch
x = torch.rand(10, 512, 512, 3)
y = x.permute(0, 3, 1, 2)

x.shape:torch.Size([10, 512, 512, 3])

y.shape:torch.Size([10, 3, 512, 512])