Pytorch torchvision.transforms中ToPILImage和ToTensor中的维度转换问题


Pytorch torchvision.transforms中ToPILImage和ToTensor中的维度转换问题

ToPILImage

官方文档中的解释
即只有是Tensor的时候才会进行维度的转变，numpy的时候则不进行转换。

ToTensor

官方文档中的解释
无论是Tensor还是ndarry都会进行维度的转换

举例

  • ndarry->
1
2
3
4
5
6
7
x = np.ones((128, 128, 3))
x_pil = transforms.ToPILImage()(np.uint8(x))   # (128, 128, 3)这里的np.uint8()是必须的
x_tensor = transforms.ToTensor()(x_pil)
print(x_tensor.size())

# out:
#torch.Size([3, 128, 128])
  • tensor->
1
2
3
4
5
6
7
8
9
x = torch.ones((3, 128, 128))
x_pil = T.ToPILImage()(x)
print(x_pil.size)
x_tensor = T.ToTensor()(x_pil)
print(x_tensor.size())

# out：
# torch.Size([128, 128, 3])
# torch.Size([3, 128, 128])