在训练神经网络的过程中最重要的工具是

  • 数据
  • 可视化
  • GPU加速

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

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

数据集可能存在两方面的问题

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

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

对PIL Image的常见操作：

  • Resize：调整图片尺寸
  • CenterCrop、RandomCrop、RandomSizedCrop：裁剪图片
  • Pad：填充
  • ToTensor：将PIL Image对象转成Tensor，会自动将[0, 255]归一化至[0, 1]

对Tensor的常见操作：

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

如果要对图片进行多个操作，可通过Compose将这些操作拼接起来:

from torchvision import transforms as T

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

Dataset只负责数据的抽象，一次调用__getitem__只返回一个样本。在训练神经网络时，是对一个batch的数据进行操作，同时还需要对数据进行shuffle和并行加速等。DataLoader可以实现这些功能。

DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, num_workers=0, collate_fn=default_collate, pin_memory=False, drop_last=False)
  • dataset：加载的数据集（Dataset对象）
  • batch_size：批大小
  • shuffle：是否将数据打乱
  • sampler：样本抽样
  • num_workers：使用多进程加载的进程数，0代表不使用多进程
  • collate_fn：如何将多个样本数据拼接成一个batch，一般使用默认的拼接方式即可
  • pin_memory：是否将数据保存在pin memory区，pin memory中的数据转到GPU会快一点
  • drop_last：dataset中的数据个数可能不是batch_size的整数倍，drop_last为True会将多出来不足一个batch的数据丢弃

torchvision主要包含以下三部分：

  • models：提供深度学习中各种经典网络的网络结构及预训练好的模型，包括AlexNet、VGG系列、ResNet系列、Inception系列等
  • datasets：提供常用的数据集加载，设计上都是继承torch.utils.data.Dataset，主要包括MNIST、CIFAR10/100、ImageNet、COCO等
  • transforms：提供常用的数据预处理操作，主要包括对Tensor及PIL Image对象的操作
from torchvision import models
from torch import nn

resnet34 = models.resnet34(pretrained=True, num_classes=1000)
resnet34.fc = nn.Linear(512, 10)

from torchvision import datasets

dataset = datasets.MNIST('data/', download=True, train=False, transform=transform)

torchvision还提供两个常用的函数：

  • make_grid：能将多张图片拼接在一个网格中
  • save_img：能将Tensor保存成图片

在PyTorch中以下数据结构分为CPU和GPU两个版本

  • Tensor
  • Variable
  • nn.Module
    它们都带有一个cuda方法，调用此方法即可将其转为对应的GPU对象。

tensor.cuda和variable.cuda都会返回一个新对象，这个新对象的数据已转至GPU，而之前的tensor/variable的数据都还在原来的设备上(CPU)。module.cuda会将所有的数据都迁移至GPU，并返回自己。所以module = module.cuda()和module.cuda()的效果相同。

在PyTorch中，以下对象可以持久化到硬盘，并能通过相应的方法加载到内存中

  • Tensor
  • Variable
  • nn.Module
  • Optimizer
from torchvision.models import AlexNet
model = AlexNet()
# module的state_dict是一个字典
model.state_dict().keys()
# Module对象的保存于加载
t.save(model.state_dict(), 'alexnet.pth')
model.load_state_dict(t.load('alexnet.pth'))