PyTorch 入门常见问题

Feb 26 2018

注意：此页面已经严重过时，不适用于 PyTorch 0.4 及之后的版本，有待修改。

最近在复现一些比较新的 Paper，里面的神经网络模型都比较多样化，训练方式也千奇百怪。对于 Caffe 和 TensorFlow 这样的预编译框架，不灵活这一缺陷就变得十分明显了。

而这方面正是 PyTorch 的强项，它是一个非常 Pythonic 的深度学习框架，一切的操作（包括模型的定义、训练、测试）都十分地符合 Python 的简单的哲学，可以非常轻松快速地构建出一个十分怪异的模型，非常适合科研人员。

作为一个去年年初才发布的框架，现在的流行程度已经快赶上 TensorFlow 了。其论坛和Github上面的开发者们都十分地活跃，对于新手非常友好。

PyTorch 的文档也非常地详细，介绍的很清楚，基本上遇到问题读一读文档都能得到一个很明确的答案。实在不行可以直接看源码，源码写的也十分优美。

因为官方社区和文档都太详细了，所以本文的价值大概仅仅是可以让初学者少踩一些坑吧。

Get Started

如何安装

官网首页就很清楚地介绍了安装方法，推荐使用 Anaconda 安装：

conda install pytorch torchvision -c pytorch

如何创建第一个神经网络

不同于 Caffe 把网络模型参数定义在一个配置文件内，以及 TensorFlow 的面向过程的搭建方式，PyTorch 里面网络模型的定义采用面向对象的方法，只需要继承 torch.nn.Module 类，并且重载里面的 __init__() 方法和 forward() 方法，便完成了一个神经网络的构造。

例如创建一个包含一个卷积层、一个池化层、一个全连接层的神经网络：

import torch
import torch.nn as nn
import torch.nn.functional as F

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv = nn.Conv2d(3, 8, 5)
        self.pool = nn.MaxPool2d(3)
        self.fc = nn.Linear(256, 10)

    def forward(self, x):
        x = self.conv(x)
        x = F.relu(x)
        x = self.pool(x)
        x = x.view(-1, 256)
        x = self.fc(x)
        return x

该网络的输入 x 是一个 4 维的 tensor，尺寸为 (b, c, h, w)。其中 b 是 batch_size，c 是 channel，图片的 channel 一般是 3，在 PyTorch 当中 channel 是在 h 和 w 前面的。h 是图片的高度，w 是图片的宽度。

nn.Conv2d(3, 8, 5)，卷积层，表示输入的 channel 是 3，输出的 channel 是 8，卷积核的尺寸是 3x3。
nn.MaxPool2d(3)，最大池化层，表示核的尺寸是 3x3。核的尺寸也可以是矩形的。
nn.Linear(256, 10)，全连接层，表示输入的维度为 256，输出的维度为 10。
F.relu(x)，ReLU 激活函数层。
x.view(-1, 256)，相当于 reshape，把一个高维的 tensor 变形为一维的向量，接入全连接层。

特别要说明的是，一个神经网络真的就这样定义好了，完全不需要再去定义 backward 函数，PyTorch 会自动计算各参数的梯度。

接下来是训练：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable

net = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)

for epoch in range(10):
    for i, data in enumerate(dataloder):
        x, label = data
        x, label = Variable(x), Variable(label)

        output = net.forward(x)
        loss = criterion(output, label)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

以上先实例化了一个模型，定义损失函数和优化器，然后分 10 个 epoch，每个 epoch 里面遍历一遍训练集。

criterion，损失函数，这里使用了交叉熵损失函数。
optimizer，优化器，这里选择了梯度下降作为是优化器，学习率是 0.01，动量是 0.9。
dataloder，数据加载器，这里就没有具体写数据是怎么来的了，其实是用 Python 怎么方便怎么来，不过一般推荐使用 torch.utils.data.DataLoader。
Variable，原始 tensor 类型的封装，为了达到自动求解梯度的效果。
optimizer.zero_grad()，将梯度清零。
loss.backward()，将误差 loss 反向传播，得到各个参数的梯度。
optimizer.step()，求得梯度之后，使用特定的优化方法（比如梯度下降）来更新参数。

于是就完成了一个完全自定义的神经网络模型从头到尾的训练过程，代码十分地简单。

如何改造已有的网络模型 / 如何提取已有模型某一层的输出

PyTorch 附带的 TorchVision 库里面包含很多经典的数据集（MNIST，CIFAR，Imagenet-12，COCO 等）和模型（AlexNet，VGG，ResNet，SqueezeNet，DenseNet，Inception 等）。经典的模型一般可以作为一个比较理想的特征提取工具。所以经常会在自己的模型里面引入一个 PreTrained 的模型作为 Backbone。然而在大多数情况下我们并不需要模型中最后几层用于分类任务的全连接层，所以还需要对模型进行改造。

下面是引入一个 PreTrained 的 ResNet50 作为 Backbone，并把最后的全连接层去掉，加入新的全连接层。

import torch
import torch.nn as nn
from torchvision.models import resnet50

def MyNet(nn.mudule):
    def __init__(self):
        super(MyNet, self).__init__()
        # 生成一个 Pretrained 的 ResNet50 模型
        resnet = resnet50(pretrained=True)
        # 获取 ResNet50 的所有 Layers
        layers = resnet.children()
        # 将最后一层去掉
        reserve_layers = list(layers)[:-1]

        # 由剩余的 Layers 重新构建新的模型
        self.new_model = nn.Sequential(*reserve_layers)
        # 自定义的全连接层，输入的维度是 2048，输出的是 10
        self.new_fc = nn.Linear(2048, 10)

    def forward(self, x):
        x = self.new_model.forward(x)
        x = self.new_fc(x)
        return x

以上就是一个改造模型的例子。若想获取一个 PreTrained 模型某一层的输出，方法与上面类似：在 __init__ 函数里面改造模型，然后在 forward 函数里面 return 你想要的输出即可。

如何在网络里面加入自定义的 Layer

有时候可能需要一些自定义的 Layers，它们不属于 Convolution、Pooling、Full-Connected 之中的任何一种，这些层既参与 forward 的计算，又会在 backward 中更新参数。

例如实现一个加权平均的 Average Pooling 层，输入尺寸是 (c, h, w) ，输出尺寸是 (c, 1, 1)，拥有的参数（即加权的权重）的尺寸为 (h, w)。

import torch
import torch.nn as nn

def WeightedPool2d(nn.module):
    def __init__(self, h, w):
        super(WeightedPool2d, self).__init__()
        init_val = torch.randn(h, w)
        # 若想实现在 backward 中更新参数，需要事先声明该变量为可训练的参数
        self.w = nn.Parameter(init_val)

    def forward(self, x):
        x = torch.mul(x, self.w)
        x = x.sum(dim=1, keepdim=True).sum(dim=2, keepdim=True)
        return x

def MyNet(nn.module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.conv = nn.Conv2d(3, 8, 5)
        # 假设上一层的输出尺寸为 (8, 128, 96)
        self.pool = WeightedPool2d(128, 96)
        self.fc = nn.Linear(8, 10)

    def forward(self, x):
        x = self.conv(x)
        x = self.pool(x)
        x = x.view(-1, 8)
        x = self.fc(x)
        return x

可以看出，实现一个自定义的层非常简单，在 forward 函数中可以很方便地使用各种自定义的运算。唯一需要注意的一点就是需要事先声明一个可训练的参数。

若是一系列的 Parameter ，则还需要声明为 ParameterList 。

Tricks

如何只学习部分模型参数 / 如何动态调整学习率

要想控制需要更新的参数的 Layers，以及每一个 Layer 的学习率，还有每一个 Epoch 的学习率，最好是用自定义的 Optimizer 来实现，因为定义一个 Optimizer 过程很简单，且几乎没什么额外开支。

以下是一个自定义的随机梯度下降的 Optimizer 的例子：

optimizer = optim.SGD([
    { 'params': net.new_model.parameters(), 'lr': 0.001 },
    { 'params': net.new_fc.parameters() }
], lr=0.01)

上面定义了一个 Optimizer，它只更新 new_model 部分和 new_fc 部分的参数（其它 Layers 的参数不更新）， new_model 部分的学习率是 0.001，其它部分的学习率是 0.01。

若想随着 Epoch 的增加来动态调整学习率，只需在每个 Epoch 中重新定义一个 Optimizer 即可。

如何加载数据集

很多时候我们需要加载自己的训练集，该训练集可能是二进制的，也可能是原始的 JPG 图片，这种情况下最好还是继承 Dataset 类比较方便。

Dataset 类的本质是定义了数据所在的位置，以及数据需要预处理的方法。至于数据的位置是在硬盘里面还是提前加载到内存里面，由该类的内部实现决定。

以下是一个加载 Market1501 数据集的例子：

import os
import re
import torch
from PIL import Image
import torch.utils.data as data

class Market1501(data.Dataset):

    base_folder = 'Market-1501-v15.09.15'
    train_folder = 'bounding_box_train'

    # root 表示数据集在硬盘中的根路径
    # transform 表示原始数据需要进行的变换（如归一化，或者扩充数据集的变换）
    def __init__(self, root, transform=None):
        self.root = root
        self.data_type = data_type
        self.transform = transform

        self.folder = os.path.join(self.root, self.base_folder, self.train_folder)
        self.pattern = re.compile(r'^(\-1|\d{4})_c(\d)s\d_\d{6}_\d{2}.*\.jpg$')
        self.file_list = list(filter(self.pattern.search, os.listdir(self.folder)))

    # 必须实现的一个函数， DataLoader 就是依靠此函数来取数据的。
    # 该实现方法为每次都到硬盘中读取原始数据并做相应的转化，比较省内存但会费时间。
    def __getitem__(self, index):
        return self.load_image(self.file_list[index])

    # 自定义数据集的 Load 方法
    def load_image(self, filename):
        label, camera = re.findall(self.pattern, filename)[0]
        label, camera = int(label), int(camera)
        img = Image.open(os.path.join(self.folder, filename))
        img.load()

        if self.transform is not None:
            img = self.transform(img)

        return img, label, camera

    # 必须实现的一个函数
    def __len__(self):
        return len(self.file_list)

如何扩充训练集

扩充数据集一般的方法就是对原始数据集做一些变换，这些操作可以集成在读取数据的过程中完成，只需定义合适的 transform 方法即可。

以下是比较常见的一些 transform 方法：

import torchvision.transforms as transforms

transform = transforms.Compose([
    # 水平移动
    transforms.RandomHorizontalFlip(),
    # 变形
    transforms.Resize((768, 256)),
    transforms.ToTensor(),
    # 正则化
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

该 transform 传递给 Dataset，之后当 DataLoader 每次去向 Dataset 取数据时，有一些变换（如随机水平移动）每次都会生成不同的图片，也就达到了扩充数据集的目的。

如何使用 GPU

只需要改动三个地方：将 Net 放到 GPU，将数据放到 GPU，将 Loss 放到 GPU。

借用上面的训练过程，只需简单地改动三行：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable

# net = Net()
net = Net().cuda()
# criterion = nn.CrossEntropyLoss()
criterion = nn.CrossEntropyLoss().cuda()
optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)

for epoch in range(10):
    for i, data in enumerate(dataloder):
        x, label = data
        # x, label = Variable(x), Variable(label)
        x, label = Variable(x).cuda(), Variable(label).cuda()

        output = net.forward(x)
        loss = criterion(output, label)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

如何并行使用 GPU（单机多卡）

使用单机多卡来训练，本质上是将训练数据随机分成多份，每一张卡 forward 其中一部分，然后再合并起来求梯度。

实现单机多卡很简单，只需要按如下方式声明网络即可：

import torch
from torch.nn import DataParallel

net = DataParallel(Net().cuda())

使用了 DataParallel 之后，相当于原始的 model 又进行了一次封装，所以原本使用 net.conv 来获取 layers 的方式需要改为 net.module.conv 。

DataParallel 只能在 GPU 模式下使用。如果使用 CPU 模式，则需要关闭 DataParallel（其实不关闭也没什么影响了）。Ref

如何使用分布式训练（多机多卡）

使用多机并行训练，本质上是将训练数据随机分成多份，每一个机器 forward 并且 backward 其中一部分，然后通过共享梯度来实现参数更新。

PyTorch 的多机多卡分布式训练一般使用 gloo 作为后端（作为平常使用的话其实可以不用在意这些细节），以下是一个简单的分布式训练的例子：

import torch
from torch.nn.parallel import DistributedDataParallel
import torch.distributed as dist

# 每个节点都要用以下语句初始化
# init_method 是 master 机器的 IP 和端口，worker 们只需要与 master 机器通信
# world_size 是结点数量
# rank 是该 worker 的序号，不同结点的 rank 是不同的
dist.init_process_group(backend='gloo', init_method=tcp://192.168.1.101:2222,
    world_size=5, rank=0)

# 定义 dataset
trainset = Market1501(root=args.dataset, data_type='train', transform=transform)
# DistributedSampler 的作用是将训练数据随机采样，送给不同的结点去 forward 和 backward
train_sampler = DistributedSampler(trainset)
# 定义 dataloader
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64,
    shuffle=(train_sampler is None), num_workers=5, pin_memory=True, sampler=train_sampler)

# 定义分布式的模型
net = Net()
net = net.cuda()
net = DistributedDataParallel(net)

for epoch in range(20):

    # 设置当前的 epoch，为了让不同的结点之间保持同步。
    train_sampler.set_epoch(epoch)

    # 以下的过程就是跟平常非分布式的一样了
    for i, data in enumerate(dataloder):
        x, label = data
        x, label = Variable(x).cuda(), Variable(label).cuda()

        output = net.forward(x)
        loss = criterion(output, label)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

训练结束后的 Bug ，这个不影响训练，但是如果是流水线工作的话就比较麻烦，训练结束之后进程就断了，还需要手动启动接下来的工作。

坑

如何保存 / 加载模型参数

最普通的方式就是先获取参数字典，然后序列化即可：

import torch
state = net.state_dict()
torch.save(state, './mynet.model')

一些常见的问题需要注意一下：

若使用 GPU 来训练，最好在保存之前转变为 CPU 模式： state = net.cpu().state_dict()
若使用单机多卡、多机多卡训练，最好把封装的 module 拿掉： state = net.module.state_dict()

如何只 forward 不 backward

若在训练好的模型上只 forward 而不 backward，最好在定义 Variable 时设置 volatile=True，避免不必要的运算和内存占用。Ref

Deep Learning , Tools