transforms.ToTensor()本身有维度转换功能

源码中的一行代码:

img = torch.from_numpy(pic.transpose((2, 0, 1)))

 再看我写的实验:

import cv2,numpy as np
from torchvision import transforms

trans=transforms.Compose(
    [
        transforms.ToTensor(),

    ])
print(np.shape(cv2.imread("test.jpg")))
print(trans(cv2.imread("test.jpg")).size())

输出结果为:

(96, 96, 3)
torch.Size([3, 96, 96])

 也就是说从opencv读到的图片是通道数在第三个维度,现在经过ToTensor操作cv2图片变成了torch image类型,也就是通道在第一个维度。接下来就不用自己显式的用tensor的permute方法来执行维度变化了。可以看我另一篇关于permute的用法:PyTorch中permute的用法

你可能感兴趣的:(pytorch学习,openCV)