深度学习框架PyTorch入门与实践：第五章 PyTorch中常用的工具

在训练神经网络的过程中需要用到很多工具，其中最重要的三部分是数据处理、可视化和GPU加速。本章主要介绍PyTorch在这几方面常用的工具，合理使用这些工具能极大地提高编程效率。

5.1 数据处理

在解决深度学习问题的过程中，往往需要花费大量的精力去处理数据，包括图像、文本、语音或其他二进制数据等。数据的处理对训练神经网络来说十分重要，良好的数据处理不仅会加速模型训练，也会提高模型效果。考虑到这一点，PyTorch提供了几个高效便捷的工具，以便使用者进行数据处理或者增强等操作，同时可通过并行化加速数据加载。

（1）数据加载

在PyTorch中，数据加载可通过自定义的数据集对象实现。数据集对象被抽象为Dataset，实现自定义的数据集需要继承Dataset，并实现两个Python魔法方法。

• getitem：返回一条数据或一个样本。obj[index]等价于obj.getitem(index)。
• len：返回样本的数量。len(obj)等价于obj.len()。

这里我们以Kaggle经典挑战赛“Dogs vs. Cats”的数据为例，详细讲解如何处理数据。“Dogs vs. Cats”是一个分类问题，判断一张图片是狗还是猫，其所有图片都存放在一个文件夹下，根据文件名的前缀判断是狗还是猫。

import torch as t
from torch.utils import data
import os
from PIL import  Image
import numpy as np

class DogCat(data.Dataset):
    def __init__(self, root):
        imgs = os.listdir(root)
        # 所有图片的绝对路径
        # 这里不实际加载图片，只是指定路径，当调用__getitem__时才会真正读图片
        self.imgs = [os.path.join(root, img) for img in imgs]
        
    def __getitem__(self, index):
        img_path = self.imgs[index]
        # dog->1， cat->0
        label = 1 if 'dog' in img_path.split('/')[-1] else 0
        pil_img = Image.open(img_path)
        array = np.asarray(pil_img)
        data = t.from_numpy(array)
        return data, label
    
    def __len__(self):
        return len(self.imgs)

dataset = DogCat('./data/dogcat/')
img, label = dataset[0] # 相当于调用dataset.__getitem__(0)
for img, label in dataset:
    print(img.size(), img.float().mean(), label)

输出：

torch.Size([500, 497, 3]) tensor(106.4915) 0
torch.Size([499, 379, 3]) tensor(171.8085) 0
torch.Size([236, 289, 3]) tensor(130.3004) 0
torch.Size([374, 499, 3]) tensor(115.5177) 0
torch.Size([375, 499, 3]) tensor(116.8139) 1
torch.Size([375, 499, 3]) tensor(150.5080) 1
torch.Size([377, 499, 3]) tensor(151.7174) 1
torch.Size([400, 300, 3]) tensor(128.1550) 1

通过上面的代码，我们学习了如何自定义自己的数据集，并可以依次获取。但这里返回的数据不适合实际使用，因其具有如下两方面问题：

• 返回样本的形状不一，因每张图片的大小不一样，这对于需要取batch训练的神经网络来说很不友好。
• 返回样本的数值较大，未归一化至[-1, 1]。

针对上述问题，PyTorch提供了torchvision。它是一个视觉工具包，提供了很多视觉图像处理的工具，其中transforms模块提供了对PIL Image对象和Tensor对象的常用操作。

对PIL Image的操作包括：

• Scale：调整图片尺寸，长宽比保持不变
• CenterCrop、`RandomCrop、RandomResizedCrop： 裁剪图片
• Pad：填充
• ToTensor：将PIL Image对象转成Tensor，会自动将[0, 255]归一化至[0, 1]

对Tensor的操作包括：

• Normalize：标准化，即减均值，除以标准差
• ToPILImage：将Tensor转为PIL Image对象

如果要对图片进行多个操作，可通过Compose函数将这些操作拼接起来，类似于nn.Sequential。注意，这些操作定义后是以函数的形式存在，真正使用时需调用它的__call__方法，这点类似于nn.Module。例如要将图片调整为224×224，首先应构建这个操作trans = Resize((224, 224))，然后调用trans(img)。下面我们就用transforms的这些操作来优化上面实现的dataset。

import os
from PIL import  Image
import numpy as np
from torchvision import transforms as T

transform = T.Compose([
    T.Resize(224), # 缩放图片(Image)，保持长宽比不变，最短边为224像素
    T.CenterCrop(224), # 从图片中间切出224*224的图片
    T.ToTensor(), # 将图片(Image)转成Tensor，归一化至[0, 1]
    T.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5]) # 标准化至[-1, 1]，规定均值和标准差
])

class DogCat(data.Dataset):
    def __init__(self, root, transforms=None):
        imgs = os.listdir(root)
        self.imgs = [os.path.join(root, img) for img in imgs]
        self.transforms=transforms
        
    def __getitem__(self, index):
        img_path = self.imgs[index]
        label = 0 if 'dog' in img_path.split('/')[-1] else 1
        data = Image.open(img_path)
        if self.transforms:
            data = self.transforms(data)
        return data, label
    
    def __len__(self):
        return len(self.imgs)

dataset = DogCat('./data/dogcat/', transforms=transform)
img, label = dataset[0]
for img, label in dataset:
    print(img.size(), label)

输出：

torch.Size([3, 224, 224]) 1
torch.Size([3, 224, 224]) 1
torch.Size([3, 224, 224]) 1
torch.Size([3, 224, 224]) 1
torch.Size([3, 224, 224]) 0
torch.Size([3, 224, 224]) 0
torch.Size([3, 224, 224]) 0
torch.Size([3, 224, 224]) 0

除了上述操作之外，transforms还可通过Lambda封装自定义的转换策略。例如想对PIL Image进行随机旋转，则可写成这样trans=T.Lambda(lambda img: img.rotate(random()*360))。

torchvision已经预先实现了常用的Dataset，包括前面使用过的CIFAR-10，以及ImageNet、COCO、MNIST、LSUN等数据集，可通过诸如torchvision.datasets.CIFAR10来调用，具体使用方法请参看官方文档。在这里介绍一个会经常使用到的Dataset——ImageFolder，它的实现和上述的DogCat很相似。ImageFolder假设所有的文件按文件夹保存，每个文件夹下存储同一个类别的图片，文件夹名为类名，其构造函数如下：

ImageFolder(root, transform=None, target_transform=None, loader=default_loader)

它主要有四个参数：

• root：在root指定的路径下寻找图片。
• transform：对PIL Image进行的转换操作，transform的输入是使用loader读取图片的返回对象。
• target_transform：对label的转换。
• loader：给定路径后如何读取图片，默认读取为RGB格式的PIL Image对象。

label是按照文件夹名顺序排序后存成字典，即{类名:类序号(从0开始)}，一般来说最好直接将文件夹命名为从0开始的数字，这样会和ImageFolder实际的label一致，如果不是这种命名规范，建议看看self.class_to_idx属性以了解label和文件夹名的映射关系。

from torchvision.datasets import ImageFolder
dataset = ImageFolder('data/dogcat_2/')

# cat文件夹的图片对应label 0，dog对应1
dataset.class_to_idx

输出：

{'cat': 0, 'dog': 1}

# 所有图片的路径和对应的label
dataset.imgs

输出：

[('data/dogcat_2/cat\\cat.12484.jpg', 0),
 ('data/dogcat_2/cat\\cat.12485.jpg', 0),
 ('data/dogcat_2/cat\\cat.12486.jpg', 0),
 ('data/dogcat_2/cat\\cat.12487.jpg', 0),
 ('data/dogcat_2/dog\\dog.12496.jpg', 1),
 ('data/dogcat_2/dog\\dog.12497.jpg', 1),
 ('data/dogcat_2/dog\\dog.12498.jpg', 1),
 ('data/dogcat_2/dog\\dog.12499.jpg', 1)]

# 没有任何的transform，所以返回的还是PIL Image对象
dataset[0][1] # 第一维是第几张图，第二维为1返回label
dataset[0][0] # 第一维是第几张图，第二维为0返回图片数据

输出：

# 加上transform
normalize = T.Normalize(mean=[0.4, 0.4, 0.4], std=[0.2, 0.2, 0.2])
transform  = T.Compose([
         T.RandomResizedCrop(224),
         T.RandomHorizontalFlip(),
         T.ToTensor(),
         normalize,
])

dataset = ImageFolder('data/dogcat_2/', transform=transform)

# 深度学习中图片数据一般保存成CxHxW，即通道数x图片高x图片宽
dataset[0][0].size()

输出：

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

to_img = T.ToPILImage()
# 0.2和0.4是标准差和均值的近似
to_img(dataset[0][0]*0.2+0.4)

输出：

Data