PyTorch常用工具（1）数据处理

前言

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

由于内容较多，本文分成了五篇文章（1）数据处理（2）预训练模型（3）TensorBoard（4）Visdom（5）CUDA与小结。

整体结构如下：

• 1 数据处理
  • 1.1 Dataset
  • 1.2 DataLoader
• 2 预训练模型
• 3 可视化工具
• 3.1 TensorBoard
• 3.2 Visdom
• 4 使用GPU加速：CUDA
• 5 小结

全文链接：

1. PyTorch中常用的工具（1）数据处理
2. PyTorch常用工具（2）预训练模型
3. PyTorch中常用的工具（3）TensorBoard
4. PyTorch中常用的工具（4）Visdom
5. PyTorch中常用的工具（5）使用GPU加速：CUDA

1 数据处理

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

1.1 Dataset

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

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

下面以Kaggle经典挑战赛"Dogs vs. Cats"的数据为例，详细讲解如何进行数据预处理。"Dogs vs. Cats"是一个分类问题，它的任务是判断一张图片是狗还是猫。在该问题中，所有图片都存放在一个文件夹下，可以根据文件名的前缀得到它们的标签值（狗或者猫）。

In: %env LS_COLORS = None 
    !tree --charset ascii data/dogcat/

Out: env: LS_COLORS=None
     data/dogcat/
     |-- cat.12484.jpg
     |-- cat.12485.jpg
     |-- cat.12486.jpg
     |-- cat.12487.jpg
     |-- dog.12496.jpg
     |-- dog.12497.jpg
     |-- dog.12498.jpg
     `-- dog.12499.jpg

     0 directories, 8 files

In: import torch as t
    from torch.utils.data import Dataset
    print(t.__version__)

Out: 1.8.0

In: import os
    from PIL import Image
    import numpy as np

    class DogCat(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.tensor(array)
            return data, label

        def __len__(self):
            return len(self.imgs)

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

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

上面的代码讲解了如何定义自己的数据集，并对数据集进行遍历。然而，这里返回的数据并不适合实际使用，主要存在以下两个问题。

• 返回样本的形状不统一，也就是每张图片的大小不一样，这对于按batch训练的神经网络来说很不友好。

• 返回样本的数值较大，没有进行归一化。

针对上述问题，PyTorch提供了torchvision工具包。torchvision是一个视觉工具包，它提供了很多视觉图像处理的工具，其中transforms模块提供了一系列数据增强的操作。本章仅对它的部分操作进行介绍，完整内容可参考官方相关文档。

仅支持PIL Image对象的常见操作如下。

• RandomChoice：在一系列transforms操作中随机执行一个操作。

• RandomOrder：以随意顺序执行一系列transforms操作。

仅支持Tensor对象的常见操作如下。

• Normalize：标准化，即减去均值，除以标准差。
• RandomErasing：随机擦除Tensor中一个矩形区域的像素。
• ConvertImageDtype：将Tensor转换为指定的类型，并进行相应的缩放。

PIL Image对象与Tensor对象相互转换的操作如下。

• ToTensor：将$H\times W\times C$​形状的PIL Image对象转换成形状为$C\times H\times W$​的Tensor，同时会自动将[0, 255]归一化至[0, 1]。
• ToPILImage：将Tensor转为PIL Image对象。

既支持PIL Image对象，又支持Tensor对象的常见操作如下。

• Resize：调整图片尺寸。

• CenterCrop、RandomCrop、RandomResizedCrop、FiveCrop： 按照不同规则对图像进行裁剪。

• RandomAffine：随机进行仿射变换，保持图像中心不变。

• RandomGrayscale：随机将图像变为灰度图。

• RandomHorizontalFlip、RandomVerticalFlip、RandomRotation：随机水平翻转、垂直翻转、旋转图像。

如果需要对图片进行多个操作，那么可以通过transforms.Compose将这些操作拼接起来，这点类似于nn.Sequential**。注意，这些操作定义后以对象的形式存在，真正使用时需要调用__call__方法，这点类似于nn.Module。**例如，要将图片的大小调整至$224\times 224$​，首先应构建操作trans = Resize((224, 224))，然后调用trans(img)。下面使用transforms的这些操作来优化上面的Dataset：

In: 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(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)

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

除了上述操作，transforms还可以通过Lambda封装自定义的转换策略。例如，如果要对PIL Image对象进行随机旋转，那么可以写成：trans = T.Lambda(lambda img: img.rotate(random() * 360))。

与torch.nn以及torch.nn.functional类似，torchvision将transforms分解为torchvision.transforms以及torchvision.transforms.functional。相比于transforms，transforms.functional为用户提供了更加灵活的操作，读者在使用时需要自己指定所有的参数。部分transforms.functional提供的操作如下，完整内容可参考官方文档。

• adjust_brightness，adjust_contrast：调整图像的亮度、对比度。
• crop，center_crop，five_crop，ten_crop：对图像按不同规则进行裁剪。
• normalize：标准化，即减均值，除以标准差。
• to_tensor：将PIL Image对象转成Tensor。

可以看出，transforms.functional中的操作与transforms十分类似。相对于transforms而言，transforms.functional可以对多个对象以相同的参数进行操作，举例说明如下：

import torchvision.transforms.functional as TF
import random

def transforms_rotate(image1, image2):
    angle = random.randint(0, 360)
    image1 = TF.rotate(image1, angle)
    image2 = TF.rotate(image2, angle)
    return image1, image2

除了对数据进行增强操作的transforms，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, is_valid_file=None)

它主要有以下五个参数。

• root：在root指定的路径下寻找图片。
• transform：对PIL Image进行相关数据增强，transform的输入是使用loader读取图片的返回对象。
• target_transform：对label的转换。
• loader：指定加载图片的函数，默认操作是读取为PIL Image对象。
• is_valid_file：获取图像路径，检查文件的有效性。

在生成数据的label时，首先按照文件夹名进行顺序排序，然后将文件夹名保存为字典，即{类名：类序号（从0开始）}。一般来说，最好直接将文件夹命名为从0开始的数字，这样会和ImageFolder实际的label一致。如果不是这种命名规范，那么建议通过self.class_to_idx属性了解label和文件夹名的映射关系。

In: !tree --charset ASCII data/dogcat_2/

Out: data/dogcat_2/
     |-- cat
     |   |-- cat.12484.jpg
     |   |-- cat.12485.jpg
     |   |-- cat.12486.jpg
     |   `-- cat.12487.jpg
     `-- dog
         |-- dog.12496.jpg
         |-- dog.12497.jpg
         |-- dog.12498.jpg
         `-- dog.12499.jpg
    
     2 directories, 8 files

In: from torchvision.datasets import ImageFolder
    dataset = ImageFolder('data/dogcat_2/')
    # cat文件夹的图片对应label 0，dog对应1
    dataset.class_to_idx

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

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

Out: [('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)]

In: # 没有任何的transforms操作，所以返回的还是PIL Image对象
    print(dataset[0][1]) # 第一维是第几张图，第二维为1返回label
    dataset[0][0] 		 # 第二维为0返回图片数据

Out: 0

In: # 加上transforms
    transform = T.Compose([
             T.RandomResizedCrop(224),
             T.RandomHorizontalFlip(), # 水平翻转
             T.ToTensor(),
             T.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5]),
    ])

In: dataset = ImageFolder('data/dogcat_2/', transform=transform)
    # 深度学习中图片数据一般保存成C×H×W，即通道数×图片高×图片宽
    dataset[0][0].size()

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

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

1.2 DataLoader

Dataset只负责数据的抽象，调用一次__getitem__返回一个样本。然而，在训练神经网络时，一次处理的对象是一个batch的数据，同时还需要对一批数据进行打乱顺序和并行加速等操作。考虑到这一点，PyTorch提供了DataLoader实现这些功能。

DataLoader的定义如下：

DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None, num_workers=0, collate_fn=None, pin_memory=False, drop_last=False, timeout=0, worker_init_fn=None, multiprocessing_context=None, generator=None, *, prefetch_factor=2, persistent_workers=False)

它主要有以下几个参数。

• dataset：加载的数据集（Dataset对象）。
• batch_size：一个batch的大小。
• shuffle：是否将数据打乱。
• sampler：样本抽样，后续会详细介绍。
• batch_sampler：与sampler类似，一次返回一个batch的索引（该参数与batch_size、shuffle、sampler和drop_last不兼容）。
• num_workers：使用多进程加载的进程数，0代表不使用多进程。
• collate_fn： 如何将多个样本数据拼接成一个batch，一般使用默认的拼接方式即可。
• pin_memory：是否将数据保存在pin memory区，pin memory中的数据转移到GPU速度更快。
• drop_last：dataset中的数据个数可能不是batch_size的整数倍，若drop_last为True，则将多出来不足一个batch的数据丢弃。
• timeout：进程读取数据的最大时间，若超时则丢弃数据。
• worker_init_fn：每个worker的初始化函数。
• prefetch_factor：每个worker预先加载的样本数。

下面举例说明DataLoader的使用方法：

In: from torch.utils.data import DataLoader
    dataloader = DataLoader(dataset, batch_size=3, shuffle=True, num_workers=0, drop_last=False)
    dataiter = iter(dataloader)
    imgs, labels = next(dataiter)
    imgs.size() # batch_size, channel, height, width

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

DataLoader是一个可迭代（iterable）对象，可以像使用迭代器一样使用它，例如：

for batch_datas, batch_labels in dataloader:
    train()
或
dataiter = iter(dataloader)
batch_datas, batch_labels = next(dataiter)

在数据处理中，有时会出现某个样本无法读取等问题，例如某张图片损坏。此时在__getitem__函数中会抛出异常，最好的解决方案是将出错的样本剔除。如果不便于处理这种情况，那么可以返回None对象，然后在Dataloader中实现自定义的collate_fn，将空对象过滤掉。注意，这种情况下DataLoader返回的一个batch的样本数目会少于batch_size。

In: class NewDogCat(DogCat): # 继承前面实现的DogCat数据集
    def __getitem__(self, index):
        try:
            # 调用父类的获取函数，即 DogCat.__getitem__(self, index)
            return super().__getitem__(index)
        except:
            return None, None

    from torch.utils.data.dataloader import default_collate # 导入默认的拼接方式
    def my_collate_fn(batch):
        '''
        batch是一个list，每个元素是dataset的返回值，形如(data, label)
        '''
        # 过滤为None的数据
        batch = [_ for _ in batch if _[0] is not None]
        if len(batch) == 0: return t.Tensor()
        return default_collate(batch) # 用默认方式拼接过滤后的batch数据

In: dataset = NewDogCat('data/dogcat_wrong/', transforms=transform)
    dataset[8]

Out: (None, None)

In: dataloader = DataLoader(dataset, 2, collate_fn=my_collate_fn, num_workers=0, shuffle=True)
    for batch_datas, batch_labels in dataloader:
        print(batch_datas.size(), batch_labels.size())

Out: torch.Size([1, 3, 224, 224]) torch.Size([1])
     torch.Size([2, 3, 224, 224]) torch.Size([2])
     torch.Size([2, 3, 224, 224]) torch.Size([2])
     torch.Size([2, 3, 224, 224]) torch.Size([2])
     torch.Size([1, 3, 224, 224]) torch.Size([1])

从上述输出中可以看出，第1个batch的batch_size为1，这是因为有一张图片损坏，无法正常返回。最后1个batch的batch_size也为1，这是因为共有9张（包括损坏的文件）图片，无法整除2（batch_size），所以最后一个batch的样本数目小于batch_size。

对于样本损坏或数据集加载异常等情况，还可以通过其他方式解决，例如随机取一张图片代替出现异常的图片：

class NewDogCat(DogCat):
    def __getitem__(self, index):
        try:
            return super().__getitem__(index)
        except:
            new_index = random.randint(0, len(self) - 1)
            return self[new_index]

相比于丢弃异常图片而言，这种做法会更好一些，它能保证每个batch的样本数目仍然是batch_size，但是在大多数情况下，最好的方式还是对数据进行彻底清洗。

DataLoader中没有太多的魔法方法，它封装了Python的标准库Multiprocessing，能够实现多进程加速，下面对DataLoader的多进程并行原理进行简要介绍。

DataLoader默认使用单进程加载数据，这样的加载方式较慢，但在系统资源有限、数据集较小能够直接加载时十分推荐。这是因为在单进程的工作模式下，若发生异常，用户在调试时能够获取更多错误信息。当数据量较大时，可以通过num_workers参数进行多进程的数据读取，多进程并行流程如下图所示。

在多进程加载数据时，每一个进程都会拷贝Dataset对象，并执行_worker_loop函数。首先，主进程生成一个batch的数据索引，并保存在队列index_queue中。然后，每个子进程执行_worker_loop函数，根据index_queue在拷贝的Dataset对象中执行__getitem__函数，获取数据。最后，每个子进程将自身获取的数据放至work_result_queue队列中，通过collate_fn处理数据，最终得到一个batch的数据data_queue。重复执行上述流程，DataLoader就实现了多进程的数据加载，更多细节读者可以参考DataLoader的相关源码。

在Dataset和DataLoader的使用方面有以下建议。

• 将高负载的操作放在__getitem__中，例如加载图片等。在多进程加载数据时，程序会并行地调用__getitem__函数，将负载高的操作放在__getitem__函数中能够实现并行加速。

• Dataset中应当尽量仅包含只读对象，避免修改任何可变对象。在多进程加载数据时，每个子进程都会拷贝Dataset对象。如果某一个进程修改了部分数据，那么在另外一个进程的拷贝中，这部分数据并不会被修改。下面是一个不好的例子：希望self.idxs返回的结果是[0,1,2,3,4,5,6,7,8]，实际上4个进程最终的self.idxs分别是[0,4,8],[1,5],[2,6],[3,7]。而dataset.idxs则是[]， 因为它并未参与迭代，并行处理的是它的四个拷贝。

class BadDataset:
    def __init__(self):
        self.idxs = [] # 取数据的次数
    def __getitem__(self, index):
        self.idxs.append(index)
        return self.idxs
    def __len__(self):
        return 9
dataset = BadDataset()
dl = t.utils.data.DataLoader(dataset, num_workers=4)
for item in dl:
    print(item) # 注意这里self.idxs的数值
print('idxs of main', dataset.idxs) # 注意这里的idxs和__getitem__返回的idxs的区别

使用Multiprocessing库时还有另外一个问题，在使用多进程加载中，如果主程序异常终止（例如使用快捷键“Ctrl+C”强行退出），那么相应的数据加载进程可能无法正常退出。虽然发现程序已经退出了，但是GPU显存和内存仍然被占用着，通过top、ps aux也能够看到已经退出的程序，这时就需要手动强行杀掉进程，建议使用如下命令：

ps x | grep <cmdline> | awk '{print $1}' | xargs kill

• ps x：获取当前用户的所有进程。
• grep <cmdline>：找到已经停止的PyTorch程序的进程，例如通过python train.py启动程序，需要写成grep 'python train.py'。
• awk '{print $1}'：获取进程的pid。
• xargs kill：杀掉进程，根据需要可能需要写成xargs kill -9强制杀掉进程。

在执行这句命令之前，建议先确认仍有未停止进程：

ps x | grep <cmdline> 

PyTorch中还单独提供了一个Sampler模块，用来对数据进行采样。常用的有随机采样器RandomSampler，当DataLoader的shuffle参数为True时，系统会自动调用这个采样器打乱数据。默认的采样器是SequentialSampler，它会按顺序一个一个进行采样。这里介绍另外一个很有用的采样方法：WeightedRandomSampler，它会根据每个样本的权重选取数据，在样本比例不均衡的问题中，可用它进行重采样。

构建WeightedRandomSampler时需提供两个参数：每个样本的权重weights、选取的样本总数num_samples以及一个可选参数replacement。权重越大的样本被选中的概率越大，待选取的样本数目一般小于全部的样本数目。replacement用于指定是否可以重复选取某一个样本，默认为True，即允许在一个epoch中重复采样某一个数据。如果设为False，那么当某一类的样本被全部选取完，但样本数目仍然未达到num_samples时，sampler不会再从该类中选择数据，此时可能导致weights参数失效。下面举例说明：

In: dataset = DogCat('data/dogcat/', transforms=transform)
    # 假设狗的图片被取出的概率是猫的概率的两倍
    # 两类图片被取出的概率与weights的绝对大小无关，只和比值有关
    weights = [2 if label == 1 else 1 for data, label in dataset]
    weights

Out: [2, 1, 1, 2, 1, 2, 1, 2]

In: from torch.utils.data.sampler import  WeightedRandomSampler
    sampler = WeightedRandomSampler(weights,\
                                    num_samples=9,\
                                    replacement=True)
    dataloader = DataLoader(dataset,\
                            batch_size=3,\
                            sampler=sampler)
    for datas, labels in dataloader:
        print(labels.tolist())

Out: [1, 1, 0]
     [0, 1, 1]
     [1, 1, 0]

可以看出，猫狗样本比例约为1：2。同时，一共只有8个样本，但是却返回了9个，说明有样本被重复返回，这就是replacement参数的作用。下面将replacement设为False：

In: sampler = WeightedRandomSampler(weights, 8, replacement=False)
    dataloader = DataLoader(dataset, batch_size=4, sampler=sampler)
    for datas, labels in dataloader:
        print(labels.tolist())

Out: [1, 0, 1, 0]
     [1, 1, 0, 0]

在replacement为False的情况下，num_samples等于dataset的样本总数。为了不重复选取，Sampler会将每个样本都返回，weight参数不再生效。

从上面的例子中可以看出Sampler在样本采样中的作用：如果指定了Sampler，那么shuffle参数不再生效，并且sampler.num_samples会覆盖dataset的实际大小，即一个epoch返回的图片总数取决于sampler.num_samples。

本小节介绍了数据加载中两个常见的操作：Dataset与DataLoader，并结合实际数据对它们的魔法方法与底层原理进行了详细介绍。数据准备与加载是神经网络训练中最基本的环节之一，读者应该熟悉其常见操作。