背景: 今天在对pytorch tensor进行可视化的时候,发现permute和cuda两个操作的先后顺序有影响。
原始代码:
正确的操作最好是先permute在cuda,或者在numpy数据的时候用transpose进行维度变换。第二种方式先cuda,在permute就有问题。
im_data = torch.from_numpy(images).type(torch.FloatTensor).permute(0, 3, 1, 2).cuda()
im_data = torch.from_numpy(images).type(torch.FloatTensor).cuda().permute(0, 3, 1, 2)
总结:
- 先permute后cuda
- 在numpy的时候先transpose在cuda