Pytorch学习笔记（四）——torchvision工具箱

一、torchvision简介

预备知识：PIL 库的使用

有不少基于 Pytorch 的工具箱都非常实用，例如处理自然语言的 torchtext，处理音频的 torchaudio 以及处理图像视频的 torchvision。

torchvision 主要包含了一些流行的数据集，模型架构和常用的图像转换功能等。本文将围绕最常用的一些 API 展开讲解。

使用 Pytorch 官网的安装命令安装 Pytorch 时，会同时安装 torchvision 这个工具箱：

conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

二、torchvision.transforms

torchvision.transforms 是一个图像转换的工具箱，不用的图像转换组件可以通过 Compose 连接从而形成一个流水线（类似于 nn.Sequential），以实现更复杂的图像转换功能。

绝大多数转换都同时支持 PIL 的 Image 对象和张量图像，只有少数转换仅支持 Image 对象或张量图像。当然我们也可以使用 ToTensor()、ToPILImage() 等工具实现 PIL Image 和张量图像之间的相互转化。

torchvision.transforms 不仅支持转换张量图像，还支持批量地转换张量图像。单个张量图像是一个具有 (C, H, W) 形状的张量，其中 C 代表通道个数，H 和 W 代表图像的高和宽。一个 batch 的张量图像是一个具有 (B, C, H, W) 形状的张量，其中 B 代表每个 batch 中的图像个数。

例如有 30 张 尺寸为 $1024\times 768$ 的 JPG 图像，若把它们作为一个 batch，则相应的张量形状为 (30, 3, 768, 1024)。这是因为 JPG 的图像模式是 RGB，即三通道，且我们在谈论分辨率时通常是 $W\times H$ 的格式。

2.1 Image 、Tensor 与 ndarray 之间的相互转化

2.1.1 ToTensor()

ToTensor() 可以将一个 PIL Image 或一个具有 (H, W, C) 形状且数值范围在 $[0,255]$之间 的 ndarray 转换成一个形状为 (C, H, W) 且数值范围在 $[0,1]$ 之间的浮点型张量。

当然，这个转换也有前提条件。对于 PIL Image，它的图像模式必须为 (L, LA, P, I, F, RGB, YCbCr, RGBA, CMYK, 1) 中的一种；对于 ndarray，它的数据类型必须为 np.uint8。

我们先尝试将 Image 转化为 Tensor。

from torchvision import transforms
from PIL import Image

img = Image.open('./pics/1.jpg')
img2tensor = transforms.ToTensor()  # 需要先实例化
img = img2tensor(img)
print(img)

相应的输出结果为：

tensor([[[0.9176, 0.9176, 0.9176,  ..., 0.9176, 0.9176, 0.9176],
         [0.9176, 0.9176, 0.9176,  ..., 0.9176, 0.9176, 0.9176],
         [0.9176, 0.9176, 0.9176,  ..., 0.9176, 0.9176, 0.9176],
         ...,
         [0.9176, 0.9176, 0.9176,  ..., 0.9176, 0.9176, 0.9176],
         [0.9176, 0.9176, 0.9176,  ..., 0.9176, 0.9176, 0.9176],
         [0.9176, 0.9176, 0.9176,  ..., 0.9176, 0.9176, 0.9176]],

        [[0.8902, 0.8902, 0.8902,  ..., 0.8902, 0.8902, 0.8902],
         [0.8902, 0.8902, 0.8902,  ..., 0.8902, 0.8902, 0.8902],
         [0.8902, 0.8902, 0.8902,  ..., 0.8902, 0.8902, 0.8902],
         ...,
         [0.8902, 0.8902, 0.8902,  ..., 0.8902, 0.8902, 0.8902],
         [0.8902, 0.8902, 0.8902,  ..., 0.8902, 0.8902, 0.8902],
         [0.8902, 0.8902, 0.8902,  ..., 0.8902, 0.8902, 0.8902]],

        [[0.8588, 0.8588, 0.8588,  ..., 0.8588, 0.8588, 0.8588],
         [0.8588, 0.8588, 0.8588,  ..., 0.8588, 0.8588, 0.8588],
         [0.8588, 0.8588, 0.8588,  ..., 0.8588, 0.8588, 0.8588],
         ...,
         [0.8588, 0.8588, 0.8588,  ..., 0.8588, 0.8588, 0.8588],
         [0.8588, 0.8588, 0.8588,  ..., 0.8588, 0.8588, 0.8588],
         [0.8588, 0.8588, 0.8588,  ..., 0.8588, 0.8588, 0.8588]]])

可以看出输出结果已经进行了归一化。

接下来我们尝试将 ndarray 转化成 Tensor。

from torchvision import transforms
import cv2

img = cv2.imread('./pics/1.jpg')
img

输出结果：

array([[[219, 227, 234],
        [219, 227, 234],
        [219, 227, 234],
        ...,
        [219, 227, 234],
        [219, 227, 234],
        [219, 227, 234]],

       [[219, 227, 234],
        [219, 227, 234],
        [219, 227, 234],
        ...,
        [219, 227, 234],
        [219, 227, 234],
        [219, 227, 234]],

       ...,

       [[219, 227, 234],
        [219, 227, 234],
        [219, 227, 234],
        ...,
        [219, 227, 234],
        [219, 227, 234],
        [219, 227, 234]],

       [[219, 227, 234],
        [219, 227, 234],
        [219, 227, 234],
        ...,
        [219, 227, 234],
        [219, 227, 234],
        [219, 227, 234]]], dtype=uint8)

可以看出数组的形状为 (H, W, C) 且还未进行归一化。

img2tensor = transforms.ToTensor()
img = img2tensor(img)
img

输出结果：

tensor([[[0.8588, 0.8588, 0.8588,  ..., 0.8588, 0.8588, 0.8588],
         [0.8588, 0.8588, 0.8588,  ..., 0.8588, 0.8588, 0.8588],
         [0.8588, 0.8588, 0.8588,  ..., 0.8588, 0.8588, 0.8588],
         ...,
         [0.8588, 0.8588, 0.8588,  ..., 0.8588, 0.8588, 0.8588],
         [0.8588, 0.8588, 0.8588,  ..., 0.8588, 0.8588, 0.8588],
         [0.8588, 0.8588, 0.8588,  ..., 0.8588, 0.8588, 0.8588]],

        [[0.8902, 0.8902, 0.8902,  ..., 0.8902, 0.8902, 0.8902],
         [0.8902, 0.8902, 0.8902,  ..., 0.8902, 0.8902, 0.8902],
         [0.8902, 0.8902, 0.8902,  ..., 0.8902, 0.8902, 0.8902],
         ...,
         [0.8902, 0.8902, 0.8902,  ..., 0.8902, 0.8902, 0.8902],
         [0.8902, 0.8902, 0.8902,  ..., 0.8902, 0.8902, 0.8902],
         [0.8902, 0.8902, 0.8902,  ..., 0.8902, 0.8902, 0.8902]],

        [[0.9176, 0.9176, 0.9176,  ..., 0.9176, 0.9176, 0.9176],
         [0.9176, 0.9176, 0.9176,  ..., 0.9176, 0.9176, 0.9176],
         [0.9176, 0.9176, 0.9176,  ..., 0.9176, 0.9176, 0.9176],
         ...,
         [0.9176, 0.9176, 0.9176,  ..., 0.9176, 0.9176, 0.9176],
         [0.9176, 0.9176, 0.9176,  ..., 0.9176, 0.9176, 0.9176],
         [0.9176, 0.9176, 0.9176,  ..., 0.9176, 0.9176, 0.9176]]])

事实上，ToTensor 中的归一化操作均是通过除以数组中的最大元进行实现的。可能有读者已经注意到，PIL Image 和 ndarray 转化成 Tensor 后内容的顺序不一致。PIL Image 转化成 Tensor 后，排列格式为 [R, G, B]，即 img[0] 代表 R 通道；而 ndarray 转化成 Tensor 后，排列格式为 [B, G, R]。

2.1.2 PILToTensor()

如果不想进行归一化，只想单纯地把 Image 对象转换为 Tensor，则可使用 PILToTensor()。

img = Image.open('./pics/1.jpg')
img2tensor = transforms.PILToTensor()
img = img2tensor(img)
img

输出结果为：

tensor([[[234, 234, 234,  ..., 234, 234, 234],
         [234, 234, 234,  ..., 234, 234, 234],
         [234, 234, 234,  ..., 234, 234, 234],
         ...,
         [234, 234, 234,  ..., 234, 234, 234],
         [234, 234, 234,  ..., 234, 234, 234],
         [234, 234, 234,  ..., 234, 234, 234]],

        [[227, 227, 227,  ..., 227, 227, 227],
         [227, 227, 227,  ..., 227, 227, 227],
         [227, 227, 227,  ..., 227, 227, 227],
         ...,
         [227, 227, 227,  ..., 227, 227, 227],
         [227, 227, 227,  ..., 227, 227, 227],
         [227, 227, 227,  ..., 227, 227, 227]],

        [[219, 219, 219,  ..., 219, 219, 219],
         [219, 219, 219,  ..., 219, 219, 219],
         [219, 219, 219,  ..., 219, 219, 219],
         ...,
         [219, 219, 219,  ..., 219, 219, 219],
         [219, 219, 219,  ..., 219, 219, 219],
         [219, 219, 219,  ..., 219, 219, 219]]], dtype=torch.uint8)

上述代码与以下等价：

img = Image.open('./pics/1.jpg')
img = torch.tensor(np.moveaxis(np.array(img), -1, 0))
img

2.1.3 ToPILImage()

我们还可以将 Tensor 或 ndarray 转化成 Image 对象。其中 Tensor 要求形状为 (C, H, W)，ndarray 要求形状为 (H, W, C)。

# 读取图像并将其转化为Tensor
img = Image.open('./pics/1.jpg')
img2tensor = transforms.ToTensor()
img = img2tensor(img)

# 将Tensor转化为PIL Image并显示
tensor2img = transforms.ToPILImage()
img = tensor2img(img)
img

上述代码最终呈现出的图片与原图完全相同。

如果使用 opencv，即

# 读取图像并将其转化为ndarray
img = cv2.imread('./pics/1.jpg')

# 将ndarray转化为PIL Image并显示
array2img = transforms.ToPILImage()
img = array2img(img)
img

则最终呈现出的图片在色差上会与原图略有区别。这是因为在 2.1.1 节中我们就提到了，将 ndarray 转化成 Tensor 后，颜色通道是按 [B, G, R] 排列的，因此需要做一个翻转操作：

# 读取图像并将其转化为ndarray
img = cv2.imread('./pics/1.jpg')
img = img[:, :, ::-1]  # 翻转通道

# 将ndarray转化为PIL Image并显示
array2img = transforms.ToPILImage()
img = array2img(img)
img

这样一来最终输出与原图保持一致。

---

更严谨地来讲，在用 opencv 读取图像之后，ndarray 就是按照 [B, G, R] 形式进行排列的，并不是 ToTensor() 导致的，读者可用下方的代码自行验证：

img = cv2.imread('./pics/1.jpg')
print(np.moveaxis(img, -1, 0))

为避免这一现象，我们可以重新改写 opencv 的读取方法：

def cv_read(path):
    import cv2
    return cv2.imread(path)[:, :, ::-1]

这样一来，在使用 ToPILImage() 进行转化时，就可以获得与原图完全一致的图像：

img = cv_read('./pics/1.jpg')
array2img = transforms.ToPILImage()
img = array2img(img)
img

2.2 常见的图像操作

导入 torchvision.transforms.functional 以实现常见的图像操作：

import torchvision.transforms.functional as TF

为方便对比，先展示原图：

2.2.1 TF.adjust_brightness()

用于调整图片亮度，控制亮度的参数必须为非负的。0 代表最暗，即返回一张全黑图，1 对应着原图的亮度。

img = Image.open('./pics/1.jpg')
TF.adjust_brightness(img, 0.5)

2.2.2 TF.adjust_contrast()

用于调整图片对比度，控制对比度的参数必须为非负的。0 代表没有对比度，即返回一张全黑图，1 对应原图。

img = Image.open('./pics/1.jpg')
TF.adjust_contrast(img, 2)

2.2.3 TF.adjust_saturation()

用于调整图片饱和度，控制饱和度的参数必须为非负的。0 代表没有饱和度，即返回一张黑白图像，1 对应原图。

img = Image.open('./pics/1.jpg')
TF.adjust_saturation(img, 3)

2.2.4 TF.adjust_sharpness()

用于调整图片锐度，控制锐度的参数必须为非负的。0 代表没有锐度，即返回一张模糊图像，1 对应原图。

img = Image.open('./pics/1.jpg')
TF.adjust_sharpness(img, 6)

2.2.5 TF.center_crop()

将图像的中间区域裁剪下来。区域参数为一个元组：(H, W)。

img = Image.open('./pics/1.jpg')
TF.center_crop(img, (320, 512))

2.2.6 TF.crop()

裁剪图像中的指定区域。

需要注意的是，TF.crop() 的坐标指定与 img.crop() 中的完全相反。在 img.crop() 中，$x$ 轴沿宽度方向；而在 TF.crop() 中，$x$ 沿高度方向。

例如裁剪左下角的 1/4 区域，该区域的左上角顶点坐标为 $(320,0)$，区域高度和宽度为 $(320,512)$，将这四个参数按序输入即可：

TF.crop(img, 320, 0, 320, 512)

2.2.7 TF.resize()

重新调整图片尺寸。尺寸参数为 (H, W)。

例如将原图调整为 $300\times 300$ 的正方形：

TF.resize(img, (300, 300))

2.2.8 TF.rotate()

将图像逆时针旋转指定角度。

例如逆时针旋转45度：

TF.rotate(img, 45)

2.2.9 TF.hflip()、TF.vflip()

这里不再展示具体图片。

TF.hflip(img)  # 水平翻转
TF.vflip(img)  # 竖直翻转

2.3 transforms.Compose()

很多时候，我们需要对大量的图片完成一系列的图像转换操作，这时候我们就能用 Compose() 将这些操作组合成一道流水线，以简化我们的代码。

例如，我们想对图片先进行水平翻转，然后调整大小至 $300\times 300$，最后将其转化为张量格式，我们可以这样操作：

trans_pipeline = transforms.Compose([
    transforms.RandomHorizontalFlip(1.0),
    transforms.Resize((300, 300)),
    transforms.ToTensor(),
])

img = Image.open('./pics/1.jpg')
img = trans_pipeline(img)

通过 ToPILImage 来显示该图像：

可以看到的确达到了我们预想的效果。

三、torchvision.io

torchvision.io 是一个用于读/写图像&视频的工具箱。本章仅介绍图片的读写。

3.1 read_image()

read_image() 用于读取 JPEG 或 PNG 格式的图片，并将其转化为 (C, H, W) 的张量，数值范围为 $[0,255]$。需要注意的是，读取后的张量是按 [R, G, B] 进行排列的。

from torchvision.io import read_image

img = read_image('./pics/1.jpg')
img

输出结果：

tensor([[[234, 234, 234,  ..., 234, 234, 234],
         [234, 234, 234,  ..., 234, 234, 234],
         [234, 234, 234,  ..., 234, 234, 234],
         ...,
         [234, 234, 234,  ..., 234, 234, 234],
         [234, 234, 234,  ..., 234, 234, 234],
         [234, 234, 234,  ..., 234, 234, 234]],

        [[227, 227, 227,  ..., 227, 227, 227],
         [227, 227, 227,  ..., 227, 227, 227],
         [227, 227, 227,  ..., 227, 227, 227],
         ...,
         [227, 227, 227,  ..., 227, 227, 227],
         [227, 227, 227,  ..., 227, 227, 227],
         [227, 227, 227,  ..., 227, 227, 227]],

        [[219, 219, 219,  ..., 219, 219, 219],
         [219, 219, 219,  ..., 219, 219, 219],
         [219, 219, 219,  ..., 219, 219, 219],
         ...,
         [219, 219, 219,  ..., 219, 219, 219],
         [219, 219, 219,  ..., 219, 219, 219],
         [219, 219, 219,  ..., 219, 219, 219]]], dtype=torch.uint8)

3.2 write_jpeg()

将形状为 (C, H, W) 的张量存储为 JPEG 格式的图片。

img = read_image('./pics/1.jpg')
write_jpeg(img, './new.jpg')

3.3 write_png()

与 3.2 同理，只不过格式换为了 PNG。

四、torchvision.datasets

torchvision.datasets 中提供了常用的图像&视频数据集，本章将主要聚焦于图像分类数据集。

from torchvision import datasets

4.1 下载数据集

以 CIFAR-10 数据集为例（官网链接），它一共包含了 $60000$ 张 $32\times 32$ 像素的图片。一共有 $10$ 类，分别是 飞机、汽车、鸟、猫、鹿、狗、青蛙、马、轮船、卡车，每一类有 $6000$ 张图片。训练集中有 $50000$ 张图片，测试集中有 $10000$ 张图片。

那么如何使用该数据集呢？我们来看下它的参数：

$$\text{datasets.CIFAR10(root,train=True,transform=None,target\_transform=None,download=False)}$$

一般情况下，本地是没有 CIFAR-10 数据集的，我们需要将其下载到本地，所以需要设置 download=True。而 root 则代表我们下载的数据集存储在何处，不妨设为 ./data/cifar10。

如果 root 中没有找到数据集，且 download=False，则会报错

RuntimeError: Dataset not found. You can use download=True to download it

如果已经下载了数据集（假如没有做任何变动的话）并且 download=True，则会显示

Files already downloaded and verified

参数 train 则是用来决定下载的是训练集还是测试集。

---

以训练集为例，我们使用如下代码下载 CIFAR-10 的训练集：

train_data = datasets.CIFAR10('./data/cifar10', train=True, download=True)

通过以下操作来进一步了解 train_data 的结构

list(train_data)
# [(, 6),
#  (, 9),
#  ...
#  (, 5),
#  ...]

len(train_data)
# 50000

train_data[0]
# (, 6)

可以看出，我们能够使用索引获取训练集中的每个图像及其对应的标签，其中图像以 PIL Image 格式存储，标签则是整型数据。

img, label = train_data[0]
img

4.2 数据集的转换

下载的数据集中的样本均是以图像形式存储的，而实际训练中我们需要张量形式，如果下载之后再一个个去用 ToTensor() 转换未免过于麻烦，我们可以直接在下载的时候就指定所需要的转换：

train_data = datasets.CIFAR10('./data/cifar10', 
							  train=True, 
							  transform=transforms.ToTensor(), 
							  download=True)

transform 参数接受任何可调用的实际参数，即我们也可以向其传递使用 Compose() 复合过后的一系列转换。

transform 是对特征的转换，target_transform 是对标签的转换。CIFAR-10 数据集的原始标签都是整型数字，因此无需进行转换。