PyTorch 口袋参考（四）
译者：飞龙
第五章：自定义 PyTorch
到目前为止，您一直在使用内置的 PyTorch 类、函数和库来设计和训练各种预定义模型、模型层和激活函数。但是，如果您有一个新颖的想法或正在进行前沿的深度学习研究怎么办？也许您发明了一个全新的层架构或激活函数。也许您开发了一个新的优化算法或一个以前从未见过的特殊损失函数。
在本章中，我将向您展示如何在 PyTorch 中创建自定义的深度学习组件和算法。我们将首先探讨如何创建自定义层和激活函数，然后看看如何将这些组件组合成自定义模型架构。接下来，我将向您展示如何创建自定义的损失函数和优化算法。最后，我们将看看如何创建用于训练、验证和测试的自定义循环。
PyTorch 提供了灵活性：您可以扩展现有库，也可以将自定义内容组合到自己的库或包中。通过创建自定义组件，您可以解决新的深度学习问题，加快训练速度，并发现执行深度学习的创新方法。
让我们开始创建一些自定义深度学习层和激活函数。
自定义层和激活
PyTorch 提供了一套广泛的内置层和激活函数。然而，PyTorch 如此受欢迎，尤其是在研究社区中，是因为创建自定义层和激活如此简单。这样做的能力可以促进实验并加速您的研究。
如果我们查看 PyTorch 源代码，我们会看到层和激活是使用功能定义和类实现创建的。功能定义指定基于输入创建输出的方式。它在nn.functional模块中定义。类实现用于创建调用此函数的对象，但它还包括从nn.Module类派生的附加功能。
例如，让我们看看全连接的nn.Linear层是如何实现的。以下代码显示了功能定义nn.functional.linear()的简化版本：
import torch
def linear(input, weight, bias=None):
if input.dim() == 2 and bias is not None:
# fused op is marginally faster
ret = torch.addmm(bias, input, weight.t())
else:
output = input.matmul(weight.t())
if bias is not None:
output += bias
ret = output
return ret
linear()函数将输入张量乘以权重矩阵，可选择添加偏置向量，并将结果返回为张量。您可以看到代码针对性能进行了优化。当输入具有两个维度且没有偏置时，应使用融合矩阵加函数torch.addmm()，因为在这种情况下速度更快。
将数学计算保留在单独的功能定义中有一个好处，即将优化与层nn.Module分开。功能定义也可以在一般编写代码时作为独立函数使用。
然而，我们通常会使用nn.Module类来对我们的神经网络进行子类化。当我们创建一个nn.Module子类时，我们获得了nn.Module对象的所有内置优势。在这种情况下，我们从nn.Module派生nn.Linear类，如下面的代码所示：
import torch.nn as nn
from torch import Tensor
class Linear(nn.Module):
def __init__(self, in_features,
out_features, bias): # ①
super(Linear, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.weight = Parameter(
torch.Tensor(out_features,
in_features))
if bias:
self.bias = Parameter(
torch.Tensor(out_features))
else:
self.register_parameter('bias', None)
self.reset_parameters()
def reset_parameters(self):
init.kaiming_uniform_(self.weight,
a=math.sqrt(5))
if self.bias is not None:
fan_in, _ = \
init._calculate_fan_in_and_fan_out(
self.weight)
bound = 1 / math.sqrt(fan_in)
init.uniform_(self.bias, -bound, bound)
def forward(self, input: Tensor) -> Tensor: # ②
return F.linear(input,
self.weight,
self.bias) # ③
①
初始化输入和输出大小、权重和偏置。
②
定义前向传递。
③
使用linear()的功能定义。
nn.Linear代码包括任何nn.Module子类所需的两种方法。一种是__init__()，它初始化类属性，即在这种情况下的输入、输出、权重和偏置。另一种是forward()方法，它定义了前向传递期间的处理。
如前面的代码所示，forward()方法经常调用与层相关的nn.functional定义。这种约定在 PyTorch 代码中经常使用于层。
创建自定义层的约定是首先创建一个实现数学运算的函数，然后创建一个nn.Module子类，该子类使用这个函数来实现层类。使用这种方法可以很容易地在 PyTorch 模型开发中尝试新的层设计。
自定义层示例（复杂线性）
接下来，我们将看看如何创建一个自定义层。在这个例子中，我们将为一种特殊类型的数字——复数创建自己的线性层。复数经常在物理学和信号处理中使用，由一对数字组成——一个“实”部分和一个“虚”部分。这两个部分都是浮点数。
PyTorch 正在添加对复杂数据类型的支持；然而，在撰写本书时，它们仍处于测试阶段。因此，我们将使用两个浮点张量来实现它们，一个用于实部，一个用于虚部。
在这种情况下，输入、权重、偏置和输出都将是复数，并且将由两个张量组成，而不是一个。复数乘法得到以下方程（其中 j 是复数 1）：
( i n r + i n i j ) ( w r + w i j ) + ( b r + b i j ) = ( i n r w r - i n i w i + b r ) + ( i n r w i + i n i w r + b i ) * j
首先，我们将创建一个复杂线性层的函数版本，如下面的代码所示：
def complex_linear(in_r, in_i, w_r, w_i, b_i, b_r):
out_r = (in_r.matmul(w_r.t())
- in_i.matmul(w_i.t()) + b_r)
out_i = (in_r.matmul(w_i.t())
- in_i.matmul(w_r.t()) + b_i)
return out_r, out_i
如你所见，该函数将复杂乘法公式应用于张量数组。接下来，我们根据nn.Module创建ComplexLinear的类版本，如下面的代码所示：
class ComplexLinear(nn.Module):
def __init__(self, in_features, out_features):
super(Linear, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.weight_r = \
Parameter(torch.randn(out_features,
in_features))
self.weight_i = \
Parameter(torch.randn(out_features,
in_features))
self.bias_r = Parameter(
torch.randn(out_features))
self.bias_i = Parameter(
torch.randn(out_features))
def forward(self, in_r, in_i):
return F.complex_linear(in_r, in_i,
self.weight_r, self.weight_i,
self.bias_r, self.bias_i)
在我们的类中，我们在__init__()函数中为实部和虚部定义了单独的权重和偏置。请注意，in_features和out_features的选项数量不会改变，因为实部和虚部的数量是相同的。我们的forward()函数只是调用我们的复杂乘法和加法操作的函数定义。
请注意，我们也可以使用 PyTorch 现有的nn.Linear层来构建我们的层，如下面的代码所示：
class ComplexLinearSimple(nn.Module):
def __init__(self, in_features, out_features):
super(ComplexLinearSimple, self).__init__()
self.fc_r = Linear(in_features,
out_features)
self.fc_i = Linear(in_features,
out_features)
def forward(self,in_r, in_i):
return (self.fc_r(in_r) - self.fc_i(in_i),
self.fc_r(in_i)+self.fc_i(in_r))
在这段代码中，我们可以免费获得nn.Linear的所有附加好处，而无需实现新的函数定义。当你创建自己的自定义层时，检查 PyTorch 的内置层，看看是否可以重用现有的类。
即使这个例子非常简单，你也可以使用相同的方法来创建更复杂的层。此外，相同的方法也可以用来创建自定义激活函数。
激活函数与 NN 层非常相似，它们通过对一组输入执行数学运算来返回输出。它们的不同之处在于，操作是逐元素执行的，并且不包括在训练过程中调整的权重和偏置等参数。因此，激活函数可以仅使用函数版本执行。
例如，让我们来看看 ReLU 激活函数。ReLU 函数对于负值为零，对于正值为线性：
def my_relu(input, thresh=0.0):
return torch.where(
input > thresh,
input,
torch.zeros_like(input))
当激活函数具有可配置参数时，通常会创建一个类版本。我们可以通过创建一个 ReLU 类来添加调整 ReLU 函数的阈值和值的功能，如下所示：
class MyReLU(nn.Module):
def __init__(self, thresh = 0.0):
super(MyReLU, self).__init__()
self.thresh = thresh
def forward(self, input):
return my_relu(input, self.thresh)
在构建 NN 时，通常使用激活函数的函数版本，但如果有的话也可以使用类版本。以下代码片段展示了如何使用torch.nn中包含的 ReLU 激活的两个版本。
这是函数版本：
import torch.nn.functional as F # ①
class SimpleNet(nn.Module):
def __init__(self, D_in, H, D_out):
super(SimpleNet, self).__init__()
self.fc1 = nn.Linear(D_in, H)
self.fc2 = nn.Linear(H, D_out)
def forward(self, x):
x = F.relu(self.fc1(x)) # ②
return self.fc2(x)
①
导入函数包的常见方式。
②
这里使用了 ReLU 的函数版本。
这是类版本：
class SimpleNet(nn.Module):
def __init__(self, D_in, H, D_out):
super(SimpleNet, self).__init__()
self.net = nn.Sequential( # ①
nn.Linear(D_in, H),
nn.ReLU(), # ②
nn.Linear(H, D_out)
)
def forward(self, x):
return self.net(x)
①
我们使用nn.Sequential()因为所有组件都是类。
②
我们正在使用 ReLU 的类版本。
自定义激活函数示例（Complex ReLU）
我们可以创建自己的自定义 ComplexReLU 激活函数来处理我们之前创建的ComplexLinear层中的复数值。以下代码展示了函数版本和类版本：
def complex_relu(in_r, in_i): # ①
return (F.relu(in_r), F.relu(in_i))
class ComplexReLU(nn.Module): # ②
def __init__(self):
super(ComplexReLU, self).__init__()
def forward(self, in_r, in_i):
return complex_relu(in_r, in_i)
①
函数版本
②
类版本
现在您已经学会了如何创建自己的层和激活函数，让我们看看如何创建自己的自定义模型架构。
自定义模型架构
在第二章和第三章中，我们使用了内置模型并从内置 PyTorch 层创建了自己的模型。在本节中，我们将探讨如何创建类似于torchvision.models的模型库，并构建灵活的模型类，根据用户提供的配置参数调整架构。
torchvision.models包提供了一个AlexNet模型类和一个alexnet()便利函数来方便其使用。让我们先看看AlexNet类：
class AlexNet(nn.Module):
def __init__(self, num_classes=1000):
super(AlexNet, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=11,
stride=4, padding=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
nn.Conv2d(64, 192, kernel_size=5,
padding=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
nn.Conv2d(192, 384, kernel_size=3,
padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(384, 256, kernel_size=3,
padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(256, 256, kernel_size=3,
padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
)
self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
self.classifier = nn.Sequential(
nn.Dropout(),
nn.Linear(256 * 6 * 6, 4096),
nn.ReLU(inplace=True),
nn.Dropout(),
nn.Linear(4096, 4096),
nn.ReLU(inplace=True),
nn.Linear(4096, num_classes),
)
def forward(self, x):
x = self.features(x)
x = self.avgpool(x)
x = torch.flatten(x, 1)
x = self.classifier(x)
return x
与所有层、激活函数和模型一样，AlexNet类派生自nn.Module类。AlexNet类是如何创建和组合子模块成为 NN 的一个很好的例子。
该库定义了三个子网络——features、avgpool和classifier。每个子网络由 PyTorch 层和激活函数组成，并按顺序连接。AlexNet 的forward()函数描述了前向传播；即输入如何被处理以形成输出。
在这种情况下，PyTorch 的torchvision.models代码提供了一个方便的函数alexnet()来实例化或创建模型并提供一些选项。这里的选项是pretrained和progress；它们确定是否加载具有预训练参数的模型以及是否显示进度条：
from torch.hub import load_state_dict_from_url
model_urls = {
'alexnet':
'https://pytorch.tips/alexnet-download',
}
def alexnet(pretrained=False,
progress=True, **kwargs):
model = AlexNet(**kwargs)
if pretrained:
state_dict = load_state_dict_from_url(
model_urls['alexnet'],
progress=progress)
model.load_state_dict(state_dict)
return model
**kwargs参数允许您向 AlexNet 模型传递其他选项。在这种情况下，您可以使用alexnet(n_classes = 10)将类别数更改为 10。该函数将使用n_classes = 10实例化 AlexNet 模型并返回模型对象。如果pretrained为True，函数将从指定的 URL 加载权重。
通过采用类似的方法，您可以创建自己的模型架构。创建一个从nn.Module派生的顶级模型。定义您的__init__()和forward()函数，并根据子网络、层和激活函数实现您的 NN。您的子网络、层和激活函数甚至可以是您自己创建的自定义的。
如您所见，nn.Module类使创建自定义模型变得容易。除了Module类外，torch.nn包还包括内置的损失函数。让我们看看如何创建自己的损失函数。
自定义损失函数
如果您回忆一下第三章，在训练 NN 模型之前，我们需要定义损失函数。损失函数或成本函数定义了我们在训练过程中希望通过调整模型权重来最小化的度量。
起初，损失函数可能看起来只是一个函数定义，但请记住，损失函数是 NN 模块参数的函数。
因此，损失函数实际上就像是一个额外的层，将 NN 的输出作为输入，并产生一个度量作为其输出。当我们进行反向传播时，我们是在损失函数上进行反向传播，而不是在 NN 上。
这使我们能够直接调用该类来计算给定 NN 输出和真实值的损失。然后我们可以一次计算所有 NN 参数的梯度，即进行反向传播。以下代码展示了如何在代码中实现这一点：
loss_fcn = nn.MSELoss() # ①
loss = loss_fcn(outputs, targets)
loss.backward()
①
有时称为criterion
首先我们实例化损失函数本身，然后调用该函数，传入输出（来自我们的模型）和目标值（来自我们的数据）。最后，我们调用backward()方法进行反向传播，并计算所有模型参数相对于损失的梯度。
与之前讨论的层类似，损失函数使用功能定义和从nn.Module类派生的类实现来实现。
mse_loss的功能定义和类实现的简化版本如下所示：
def mse_loss(input, target):
return ((input-target)**2).mean()
class MSELoss(nn.Module):
def __init__(self):
super(MSELoss, self).__init__()
def forward(self, input, target):
return F.mse_loss(input, target)
让我们创建自己的损失函数，复数的 MSE 损失。为了创建自定义损失函数，我们首先定义一个数学上描述损失函数的功能定义。然后我们将创建损失函数类，如下所示：
def complex_mse_loss(input_r, input_i,
target_r, target_i):
return (((input_r-target_r)**2).mean(),
((input_i-target_i)**2).mean())
class ComplexMSELoss(nn.Module):
def __init__(self, real_only=False):
super(ComplexMSELoss, self).__init__()
self.real_only = real_only
def forward(self, input_r, input_i,
target_r, target_i):
if (self.real_only):
return F.mse_loss(input_r, target_r)
else:
return complex_mse_loss(
input_r, input_i,
target_r, target_i)
这次，我们在类中创建了一个名为real_only的可选设置。当我们使用real_only = True实例化损失函数时，将使用mse_loss()函数而不是complex_mse_loss()函数。
正如您所看到的，PyTorch 在构建自定义模型架构和损失函数方面提供了出色的灵活性。在进行训练之前，还有一个函数可以自定义：优化器。让我们看看如何创建自定义优化器。
自定义优化器算法
优化器在训练 NN 模型中起着重要作用。优化器是在训练过程中更新模型参数的算法。当我们使用loss.backward()进行反向传播时，我们确定参数应该增加还是减少以最小化损失。优化器使用梯度来确定在每一步中参数应该改变多少并进行相应的更改。
PyTorch 有自己的子模块称为torch.optim，其中包含许多内置的优化器算法，正如我们在第三章中所看到的。要创建一个优化器，我们传入我们模型的参数和任何特定于优化器的选项。例如，以下代码创建了一个学习率为 0.01 和动量值为 0.9 的 SGD 优化器：
from torch import optim
optimizer = optim.SGD(model.parameters(),
lr=0.01, momentum=0.9)
在 PyTorch 中，我们还可以为不同的参数指定不同的选项。当您想要为模型的不同层指定不同的学习率时，这是很有用的。每组参数称为参数组。我们可以使用字典指定不同的选项，如下所示：
optim.SGD([
{'params':
model.features.parameters()},
{'params':
model.classifier.parameters(),
'lr': 1e-3}
], lr=1e-2, momentum=0.9)
假设我们正在使用 AlexNet 模型，上述代码将分类器层的学习率设置为1e-3，并使用默认学习率1e-2来训练特征层。
PyTorch 提供了一个torch.optim.Optimizer基类，以便轻松创建自定义优化器。以下是Optimizer基类的简化版本：
from collections import defaultdict
class Optimizer(object):
def __init__(self, params, defaults):
self.defaults = defaults
self.state = defaultdict(dict) # ①
self.param_groups = [] # ②
param_groups = list(params)
if len(param_groups) == 0:
raise ValueError(
"""optimizer got an
empty parameter list""")
if not isinstance(param_groups[0], dict):
param_groups = [{'params': param_groups}]
for param_group in param_groups:
self.add_param_group(param_group)
def __getstate__(self):
return {
'defaults': self.defaults,
'state': self.state,
'param_groups': self.param_groups,
}
def __setstate__(self, state):
self.__dict__.update(state)
def zero_grad(self): # ③
for group in self.param_groups:
for p in group['params']:
if p.grad is not None:
p.grad.detach_()
p.grad.zero_()
def step(self, closure): # ④
raise NotImplementedError
①
根据需要定义state。
②
根据需要定义 param_groups。
③
根据需要定义 zero_grad()。
④
您需要编写自己的 step()。
优化器有两个主要属性或组件：state 和 param_groups。state 属性是一个字典，可以在不同的优化器之间变化。它主要用于在每次调用 step() 函数之间维护值。param_groups 属性也是一个字典。它包含参数本身以及每个组的相关选项。
Optimizer 基类中的重要方法是 zero_grad() 和 step()。zero_grad() 方法用于在每次训练迭代期间将梯度归零或重置。step() 方法用于执行优化器算法，计算每个参数的变化，并更新模型对象中的参数。zero_grad() 方法已经为您实现。但是，当创建自定义优化器时，您必须创建自己的 step() 方法。
让我们通过创建我们自己简单版本的 SGD 来演示这个过程。我们的 SDG 优化器将有一个选项——学习率（LR）。在每个优化器步骤中，我们将梯度乘以 LR，并将其添加到参数中（即，调整模型的权重）：
from torch.optim import Optimizer
class SimpleSGD(Optimizer):
def __init__(self, params, lr='required'):
if lr is not 'required' and lr < 0.0:
raise ValueError(
"Invalid learning rate: {}".format(lr))
defaults = dict(lr=lr)
super(SimpleSGD, self).__init__(
params, defaults)
def step(self):
for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue
d_p = p.grad
p.add_(d_p, alpha=-group['lr'])
return
__init__() 函数设置默认选项值，并根据输入参数初始化参数组。请注意，我们不必编写任何代码来执行此操作，因为 super(SGD, self).init(params, defaults) 调用基类初始化方法。我们真正需要做的是编写 step() 方法。对于每个参数组，我们首先将参数乘以组的 LR，然后从参数本身减去该乘积。这通过调用 p.add_(d_p, alpha=-group['lr']) 完成。
以下是我们如何使用新优化器的示例：
optimizer = SimpleSGD(model.parameters(),
lr=0.001)
我们还可以使用以下代码为模型的不同层定义不同的学习率。在这里，我们假设再次使用 AlexNet 作为模型，其中包含名为 feature 和 classifier 的层：
optimizer = SimpleSGD([
{'params':
model.features.parameters()},
{'params':
model.classifier.parameters(),
'lr': 1e-3}
], lr=1e-2)
现在您可以为训练模型创建自己的优化器了，让我们看看如何创建自己的自定义训练、验证和测试循环。
自定义训练、验证和测试循环
在整本书中，我们一直在使用自定义训练、验证和测试循环。这是因为在 PyTorch 中，所有训练、验证和测试循环都是由程序员手动创建的。
与 Keras 不同，没有 fit() 或 eval() 方法来执行循环。相反，PyTorch 要求您编写自己的循环。在许多情况下，这实际上是一个好处，因为您希望控制训练过程中发生的事情。
实际上，在“生成学习—使用 DCGAN 生成 Fashion-MNIST 图像”中的参考设计演示了如何创建更复杂的训练循环。
在本节中，我们将探讨编写循环的传统方式，并讨论开发人员定制循环的常见方式。让我们回顾一些用于训练、验证和测试循环的常用代码：
for epoch in range(n_epochs):
# Training
for data in train_dataloader:
input, targets = data
optimizer.zero_grad()
output = model(input)
train_loss = criterion(output, targets)
train_loss.backward()
optimizer.step()
# Validation
with torch.no_grad():
for input, targets in val_dataloader:
output = model(input)
val_loss = criterion(output, targets)
# Testing
with torch.no_grad():
for input, targets in test_dataloader:
output = model(input)
test_loss = criterion(output, targets)
这段代码应该看起来很熟悉，因为我们在整本书中经常使用它。我们假设 n_epochs、model、criterion、optimizer、train_、val_ 和 test_dataloader 已经定义。对于每个时期，我们执行训练和验证循环。训练循环逐批次处理每个批次，将批次输入通过模型，并计算损失。然后我们执行反向传播来计算梯度，并执行优化器来更新模型的参数。
验证循环禁用梯度计算，并逐批次将验证数据通过网络传递。测试循环逐批次将测试数据通过模型传递，并计算测试数据的损失。
让我们为我们的循环添加一些额外的功能。可能性是无限的，但这个示例将演示一些简单的任务，比如打印信息、重新配置模型以及在训练过程中调整超参数。让我们走一遍以下代码，看看如何实现这一点：
for epoch in range(n_epochs):
total_train_loss = 0.0 # ①
total_val_loss = 0.0 # ①
if (epoch == epoch//2):
optimizer = optim.SGD(model.parameters(),
lr=0.001) # ②
# Training
model.train() # ③
for data in train_dataloader:
input, targets = data
optimizer.zero_grad()
output = model(input)
train_loss = criterion(output, targets)
train_loss.backward()
optimizer.step()
total_train_loss += train_loss # ①
# Validation
model.eval() # ③
with torch.no_grad():
for input, targets in val_dataloader:
output = model(input)
val_loss = criterion(output, targets)
total_val_loss += val_loss # ①
print("""Epoch: {}
Train Loss: {}
Val Loss {}""".format(
epoch, total_train_loss,
total_val_loss)) # ①
# Testing
model.eval()
with torch.no_grad():
for input, targets in test_dataloader:
output = model(input)
test_loss = criterion(output, targets)
①
打印 epoch、训练和验证损失的示例
②
重新配置模型的示例（最佳实践）
③
修改训练过程中的超参数示例
在上述代码中，我们添加了一些变量来跟踪运行的训练和验证损失，并在每个 epoch 打印它们。接下来，我们使用 train() 或 eval() 方法来配置模型进行训练或评估。这仅适用于模型的 forward() 函数在训练和评估时表现不同的情况。
例如，一些模型可能在训练过程中使用 dropout，但在验证或测试过程中不应用 dropout。在这种情况下，我们可以通过调用 model.train() 或 model.eval() 来重新配置模型，然后执行它。
最后，我们在训练过程中修改了优化器中的学习率。这使我们能够在一半的 epoch 训练后以更快的速度训练，同时在微调参数更新之后。
这个示例是如何自定义您的循环的简单演示。训练、验证和测试循环可能会更加复杂，因为您同时训练多个网络、使用多模态数据，或设计更复杂的网络，甚至可以训练其他网络。PyTorch 提供了设计用于训练、验证和测试的特殊和创新过程的灵活性。
提示
PyTorch Lightning 是一个第三方 PyTorch 包，提供了用于训练、验证和测试循环的样板模板。该包提供了一个框架，允许您创建自定义循环，而无需为每个模型实现重复输入样板代码。我们将在第八章中讨论 PyTorch Lightning。您也可以在PyTorch Lightning 网站上找到更多信息。
在本章中，您学习了如何为在 PyTorch 中开发深度学习模型创建自定义组件。随着您的模型变得越来越复杂，您可能会发现您需要训练模型的时间可能会变得相当长——也许是几天甚至几周。在下一章中，您将看到如何利用内置的 PyTorch 能力来加速和优化您的训练过程，从而显著减少整体模型开发时间。
第六章：PyTorch 加速和优化
在前几章中，您学习了如何使用 PyTorch 的内置功能，并通过创建自己的自定义组件来扩展这些功能，从而使您能够快速设计新模型和算法来训练它们。
然而，当处理非常大的数据集或更复杂的模型时，将模型训练在单个 CPU 或 GPU 上可能需要非常长的时间——可能需要几天甚至几周才能获得初步结果。训练时间更长可能会变得令人沮丧，特别是当您想要使用不同的超参数配置进行许多实验时。
在本章中，我们将探讨使用 PyTorch 加速和优化模型开发的最新技术。首先，我们将看看如何使用张量处理单元（TPU）而不是 GPU 设备，并考虑在使用 TPU 时可以提高性能的情况。接下来，我将向您展示如何使用 PyTorch 的内置功能进行并行处理和分布式训练。这将为跨多个 GPU 和多台机器训练模型提供一个快速参考，以便在更多硬件资源可用时快速扩展您的训练。在探索加速训练的方法之后，我们将看看如何使用高级技术（如超参数调整、量化和剪枝）来优化您的模型。
本章还将提供参考代码，以便轻松入门，并提供我们使用的关键软件包和库的参考资料。一旦您创建了自己的模型和训练循环，您可以返回到本章获取有关如何加速和优化训练过程的提示。
让我们开始探讨如何在 TPU 上运行您的模型。
在 TPU 上的 PyTorch
随着深度学习和人工智能的不断部署，公司正在开发定制硬件芯片或 ASIC，旨在优化硬件中的模型性能。谷歌开发了自己的用于神经网络加速的 ASIC，称为 TPU。由于 TPU 是为神经网络设计的，它没有 GPU 的一些缺点，GPU 是为图形处理而设计的。谷歌的 TPU 现在可以作为谷歌云 TPU 的一部分供您使用。您还可以在 Google Colab 上运行 TPU。
在前几章中，我向您展示了如何使用 GPU 测试和训练您的深度模型。如果您的用例符合以下条件，您应该继续使用 CPU 和 GPU 进行训练：
-
您有小型或中型模型以及小批量大小。
-
您的模型训练时间不长。
-
数据的进出是您的主要瓶颈。
-
您的计算经常是分支的或主要是逐元素完成的，或者您使用稀疏内存访问。
-
您需要使用高精度。双精度不适合在 TPU 上使用。
另一方面，有几个原因可能会导致您希望使用 TPU 而不是 GPU 进行训练。TPU 在执行密集向量和矩阵计算方面非常快速。它们针对特定工作负载进行了优化。如果您的用例符合以下情况，您应该强烈考虑使用 TPU：
-
您的模型主要由矩阵计算组成。
-
您的模型训练时间很长。
-
您希望在 TPU 上运行整个训练循环的多次迭代。
在 TPU 上运行与在 CPU 或 GPU 上运行非常相似。让我们回顾一下如何在 GPU 上训练模型的以下代码：
device = torch.device("cuda" if
torch.cuda.is_available() else "cpu") # ①
model.to(device) # ②
for epoch in range(n_epochs):
for data in trainloader:
input, labels = data
input = input.to(device) # ③
labels = labels.to(device) # ③
optimizer.zero_grad()
output = model(input)
loss = criterion(input, labels)
loss.backward()
optimizer.step()
①
如果有 GPU 可用，请将设备配置为 GPU。
②
将模型发送到设备。
③
将输入和标签发送到 GPU。
换句话说，我们将模型、输入和标签移至 GPU，其余工作由系统完成。在 TPU 上训练网络几乎与在 GPU 上训练相同，只是您需要使用PyTorch/XLA（加速线性代数）包，因为 TPU 目前不受 PyTorch 原生支持。
让我们在 Google Colab 上使用 Cloud TPU 训练我们的模型。打开一个新的 Colab 笔记本，并从运行时菜单中选择更改运行时类型。然后从“硬件加速器”下拉菜单中选择 TPU，如图 6-1 所示。Google Colab 提供免费的 Cloud TPU 系统，包括远程 CPU 主机和每个具有两个核心的四个 TPU 芯片。
图 6-1。在 Google Colab 中使用 TPU
由于 Colab 默认未安装 PyTorch/XLA，我们需要首先安装它，使用以下命令。这将安装最新的“夜间”版本，但如果需要，您可以选择其他版本：
!curl 'https://raw.githubusercontent.com/pytorch' \
'/xla/master/contrib/scripts/env-setup.py' \
-o pytorch-xla-env-setup.py
!python pytorch-xla-env-setup.py --version "nightly" # ①
<1>这些是打算在笔记本中运行的命令。在命令行上运行时，请省略“!”。
安装 PyTorch/XLA 后，我们可以导入该软件包并将数据移动到 TPU：
import torch_xla.core.xla_model as xm
device = xm.xla_device()
请注意，我们这里不使用torch.cuda.is_available()，因为它仅适用于 GPU。不幸的是，TPU 没有is_available()方法。如果您的环境未配置为 TPU，您将收到错误消息。
设备设置完成后，其余代码完全相同：
model.to(device)
for epoch in range(n_epochs):
for data in trainloader:
input, labels = data
input = input.to(device)
labels = labels.to(device)
optimizer.zero_grad()
output = model(input)
loss = criterion(input, labels)
loss.backward()
optimizer.step()
print(output.device) # ①
# out: xla:1
①
如果 Colab 配置为 TPU，您应该看到 xla:1。
PyTorch/XLA 是一个用于 XLA 操作的通用库，可能支持除 TPU 之外的其他专用 ASIC。有关 PyTorch/XLA 的更多信息，请访问PyTorch/XLA GitHub 存储库。
在 TPU 上运行仍然存在许多限制，GPU 支持更加普遍。因此，大多数 PyTorch 开发人员将首先使用单个 GPU 对其代码进行基准测试，然后再探索使用单个 TPU 或多个 GPU 加速其代码。
我们已经在本书的前面部分介绍了如何使用单个 GPU。在下一节中，我将向您展示如何在具有多个 GPU 的机器上训练您的模型。
在多个 GPU 上的 PyTorch（单台机器）
在加速训练和开发时，充分利用您可用的硬件资源非常重要。如果您有一台本地计算机或网络服务器可以访问多个 GPU，本节将向您展示如何充分利用系统上的 GPU。此外，您可能希望通过在单个实例上使用云 GPU 来扩展 GPU 资源。这通常是在考虑分布式训练方法之前的第一级扩展。
在多个 GPU 上运行代码通常称为并行处理。并行处理有两种方法：数据并行处理和模型并行处理。在数据并行处理期间，数据批次在多个 GPU 之间分割，而每个 GPU 运行模型的副本。在模型并行处理期间，模型在多个 GPU 之间分割，数据批次被管道传送到每个部分。
数据并行处理在实践中更常用。模型并行处理通常保留用于模型不适合单个 GPU 的情况。我将在本节中向您展示如何执行这两种类型的处理。
数据并行处理
图 6-2 说明了数据并行处理的工作原理。在此过程中，每个数据批次被分成N部分（N是主机上可用的 GPU 数量）。N通常是 2 的幂。每个 GPU 持有模型的副本，并且为批次的每个部分计算梯度和损失。在每次迭代结束时，梯度和损失被合并。这种方法适用于较大的批次大小和模型适合单个 GPU 的用例。
PyTorch 可以使用单进程，多线程方法或使用多进程方法来实现数据并行处理。单进程，多线程方法只需要一行额外的代码，但在许多情况下性能不佳。
图 6-2。数据并行处理
不幸的是，由于 Python 的全局解释器锁（GIL）在线程之间的争用、模型的每次迭代复制以及输入散布和输出收集引入的额外开销，多线程性能较差。您可能想尝试这种方法，因为它非常简单，但在大多数情况下，您可能会使用多进程方法。
使用 nn.DataParallel 的多线程方法
PyTorch 的nn模块原生支持多线程的数据并行处理。您只需要在将模型发送到 GPU 之前将其包装在nn.DataParallel中，如下面的代码所示。在这里，我们假设您已经实例化了您的模型：
if torch.cuda.device_count() > 1:
print("This machine has",
torch.cuda.device_count(),
"GPUs available.")
model = nn.DataParallel(model)
model.to("cuda")
首先，我们检查确保我们有多个 GPU，然后我们使用nn.DataParallel()在将模型发送到 GPU 之前设置数据并行处理。
这种多线程方法是在多个 GPU 上运行的最简单方式；然而，多进程方法通常在单台机器上表现更好。此外，多进程方法也可以用于跨多台机器运行，我们将在本章后面看到。
使用 DDP 的多进程方法（首选）
最好使用多进程方法在多个 GPU 上训练您的模型。PyTorch 通过其nn.parallel.DistributedDataProcessing模块支持这一点。分布式数据处理（DDP）可以在单台机器上的多个进程或跨多台机器的多个进程中使用。我们将从单台机器开始。
有四个步骤需要修改您的代码：
-
使用torch.distributed初始化一个进程组。
-
使用*torch.nn.to()*创建一个本地模型。
-
使用torch.nn.parallel将模型包装在 DDP 中。
-
使用torch.multiprocessing生成进程。
以下代码演示了如何将您的模型转换为 DDP 训练。我们将其分解为步骤。首先，导入必要的库：
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel \
import DistributedDataParallel as DDP
请注意，我们正在使用三个新库—torch.distributed、torch.multiprocessing和torch.nn.parallel。以下代码向您展示如何创建一个分布式训练循环：
def dist_training_loop(rank,
world_size,
dataloader,
model,
loss_fn,
optimizer):
dist.init_process_group("gloo",
rank=rank,
world_size=world_size) # ①
model = model.to(rank) # ②
ddp_model = DDP(model,
device_ids=[rank]) # ③
optimizer = optimizer(
ddp_model.parameters(),
lr=0.001)
for epochs in range(n_epochs):
for input, labels in dataloader:
input = input.to(rank)
labels = labels.to(rank) # ④
optimizer.zero_grad()
outputs = ddp_model(input) # ⑤
loss = loss_fn(outputs, labels)
loss.backward()
optimizer.step()
dist.destroy_process_group()
①
使用world_size进程设置一个进程组。
②
将模型移动到 ID 为rank的 GPU。
③
将模型包装在 DDP 中。
④
将输入和标签移动到 ID 为rank的 GPU。
⑤
调用 DDP 模型进行前向传递。
DDP 将模型状态从rank0进程广播到所有其他进程，因此我们不必担心不同进程具有具有不同初始化权重的模型。
DDP 处理了低级别的进程间通信，使您可以将模型视为本地模型。在反向传播过程中，当loss.backward()返回时，DDP 会自动同步梯度并将同步的梯度张量放在params.grad中。
现在我们已经定义了进程，我们需要使用spawn()函数创建这些进程，如下面的代码所示：
if __name__=="__main__":
world_size = 2
mp.spawn(dist_training_loop,
args=(world_size,),
nprocs=world_size,
join=True)
在这里，我们将代码作为main运行，生成两个进程，每个进程都有自己的 GPU。这就是如何在单台机器上的多个 GPU 上运行数据并行处理。
警告
GPU 设备不能在进程之间共享。
如果您的模型不适合单个 GPU 或者使用较小的批量大小，您可以考虑使用模型并行处理而不是数据并行处理。接下来我们将看看这个。
模型并行处理
图 6-3 展示了模型并行处理的工作原理。在这个过程中，模型被分割到同一台机器上的 N 个 GPU 中。如果我们按顺序处理数据批次，下一个 GPU 将始终等待前一个 GPU 完成，这违背了并行处理的目的。因此，我们需要对数据处理进行流水线处理，以便每个 GPU 在任何给定时刻都在运行。当我们对数据进行流水线处理时，只有前 N 个批次按顺序运行，然后每个后续运行会激活所有 GPU。
图 6-3. 模型并行处理
实现模型并行处理并不像数据并行处理那样简单，它需要您重新编写模型。您需要定义模型如何跨多个 GPU 分割以及数据在前向传递中如何进行流水线处理。通常通过为模型编写一个子类，具有特定数量的 GPU 的多 GPU 实现来完成这一点。
以下代码演示了 AlexNet 的双 GPU 实现：
class TwoGPUAlexNet(AlexNet):
def __init__(self):
super(ModelParallelAlexNet, self).__init__(
num_classes=num_classes,
*args,
**kwargs)
self.features.to('cuda:0')
self.avgpool.to('cuda:0')
self.classifier.to('cuda:1')
self.split_size = split_size
def forward(self, x):
splits = iter(x.split(self.split_size,
dim=0))
s_next = next(splits)
s_prev = self.seq1(s_next).to('cuda:1')
ret = []
for s_next in splits:
s_prev = self.seq2(s_prev) # ①
ret.append(self.fc(
s_prev.view(s_prev.size(0), -1)))
s_prev = self.seq1(s_next).to('cuda:1') # ②
s_prev = self.seq2(s_prev)
ret.append(self.fc(
s_prev.view(s_prev.size(0), -1)))
return torch.cat(ret)
①
s_prev 在 cuda:1 上运行。
②
s_next 在 cuda:0 上运行，可以与 s_prev 并行运行。
因为我们从 AlexNet 类派生一个子类，我们继承了它的模型结构，所以不需要创建我们自己的层。相反，我们需要描述模型的哪些部分放在 GPU0 上，哪些部分放在 GPU1 上。然后我们需要在 forward() 方法中通过每个 GPU 管道传递数据来实现 GPU 流水线。当训练模型时，您需要将标签放在最后一个 GPU 上，如下面的代码所示：
model = TwoGPUAlexNet()
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)
for epochs in range(n_epochs):
for input, labels in dataloader;
input = input.to("cuda:0")
labels = labels.to("cuda:1") # ①
optimizer.zero_grad()
outputs = model(input)
loss_fn(outputs, labels).backward()
optimizer.step()
①
将输入发送到 GPU0，将标签发送到 GPU1。
如您所见，训练循环需要更改一行代码，以确保标签位于最后一个 GPU 上，因为在计算损失之前输出将位于那里。
数据并行处理和模型并行处理是利用多个 GPU 进行加速训练的两种有效范式。如果我们能够将这两种方法结合起来并取得更好的结果，那将是多么美妙呢？让我们看看如何实现结合的方法。
结合数据并行处理和模型并行处理
您可以将数据并行处理与模型并行处理结合起来，以进一步提高性能。在这种情况下，您将使用 DDP 包装您的模型，将数据批次分发给多个进程。每个进程将使用多个 GPU，并且您的模型将被分割到每个 GPU 中。
我们只需要做两个更改。
-
将我们的多 GPU 模型类更改为接受设备作为输入。
-
在前向传递期间省略设置输出设备。DDP 将确定输入和输出数据的放置位置。
以下代码显示了如何修改多 GPU 模型：
class Simple2GPUModel(nn.Module):
def __init__(self, dev0, dev1):
super(Simple2GPUModel,
self).__init__()
self.dev0 = dev0
self.dev1 = dev1
self.net1 = torch.nn.Linear(
10, 10).to(dev0)
self.relu = torch.nn.ReLU()
self.net2 = torch.nn.Linear(
10, 5).to(dev1)
def forward(self, x):
x = x.to(self.dev0)
x = self.relu(self.net1(x))
x = x.to(self.dev1)
return self.net2(x)
在 __init__() 构造函数中，我们传入 GPU 设备对象 dev0 和 dev1，并描述模型的哪些部分位于哪些 GPU 中。这使我们能够在不同进程上实例化新模型，每个模型都有两个 GPU。forward() 方法在模型的适当位置将数据从一个 GPU 移动到下一个 GPU。
以下代码显示了训练循环的更改：
def model_parallel_training(rank, world_size):
print(f"Running DDP with a model parallel")
setup(rank, world_size)
# set up mp_model and devices for this process
dev0 = rank * 2
dev1 = rank * 2 + 1
mp_model = Simple2GPUModel(dev0, dev1)
ddp_mp_model = DDP(mp_model) # ①
loss_fn = nn.MSELoss()
optimizer = optim.SGD(
ddp_mp_model.parameters(), lr=0.001)
for epochs in range(n_epochs):
for input, labels in dataloader:
input = input.to(dev0),
labels = labels,to(dev1) # ②
optimizer.zero_grad()
outputs = ddp_mp_model(input) # ③
loss = loss_fn(outputs, labels)
loss.backward()
optimizer.step()
cleanup()
①
将模型包装在 DDP 中。
②
将输入和标签移动到适当的设备 ID。
③
输出在 dev1 上。
总之，当在多个 GPU 上使用 PyTorch 时，您有几个选项。您可以使用本节中的参考代码来实现数据并行、模型并行或组合并行处理，以加速模型训练和推断。到目前为止，我们只讨论了单台机器或云实例上的多个 GPU。
在许多情况下，在单台机器上的多个 GPU 上进行并行处理可以将训练时间减少一半或更多，您只需要升级 GPU 卡或利用更大的云 GPU 实例。但是，如果您正在训练非常复杂的模型或使用极其大型的数据集，您可能希望使用多台机器或云实例来加速训练。
好消息是，在多台机器上使用 DDP 与在单台机器上使用 DDP 并没有太大的区别。下一节将展示如何实现这一点。
分布式训练（多台机器）
如果在单台机器上训练您的 NN 模型不能满足您的需求，并且您可以访问一组服务器集群，您可以使用 PyTorch 的分布式处理能力将训练扩展到多台机器。PyTorch 的分布式子包 torch.distributed 提供了丰富的功能集，以适应各种训练架构和硬件平台。
torch.distributed 子包由三个组件组成：DDP、基于 RPC 的分布式训练（RPC）和集体通信（c10d）。我们在上一节中使用了 DDP 在单台机器上运行多个进程，它最适合数据并行处理范式。RPC 是为支持更一般的训练架构而创建的，并且可以用于除数据并行处理范式之外的分布式架构。
c10d 组件是一个用于在进程之间传输张量的通信库。c10d 被 DDP 和 RPC 组件用作后端，PyTorch 提供了 c10d API，因此您可以在自定义分布式应用中使用它。
在本书中，我们将重点介绍使用 DDP 进行分布式训练。但是，如果您有更高级的用例，您可能希望使用 RPC 或 c10d。您可以通过阅读 PyTorch 文档 了解更多信息。
对于使用 DDP 进行分布式训练，我们将遵循与在单台机器上使用多个进程相同的 DDP 过程。但是，在这种情况下，我们将在单独的机器或实例上运行每个进程。
要在多台机器上运行，我们使用一个指定配置的启动脚本来运行 DDP。启动脚本包含在 torch.distributed 中，并且可以按照以下代码执行。假设您有两个节点，节点 0 和节点 1。节点 0 是主节点，IP 地址为 192.168.1.1，空闲端口为 1234。在节点 0 上，您将运行以下脚本：
>>> python -m torch.distributed.launch
--nproc_per_node=NUM_GPUS
--nnodes=2
--node_rank=0 # ①
--master_addr="192.168.1.1"
--master_port=1234
TRAINING_SCRIPT.py (--arg1 --arg2 --arg3)
①
node_rank 被设置为节点 0。
在节点 1 上，您将运行下一个脚本。请注意，此节点的等级是 1：
>>> python -m torch.distributed.launch
--nproc_per_node=NUM_GPUS
--nnodes=2
--node_rank=1 # ①
--master_addr="192.168.1.1"
--master_port=1234
TRAINING_SCRIPT.py (--arg1 --arg2 --arg3)
①
node_rank 被设置为节点 1。
如果您想探索此脚本中的可选参数，请运行以下命令：
>>> python -m torch.distributed.launch --help
请记住，如果您不使用 DDP 范式，您应该考虑为您的用例使用 RPC 或 c10d API。并行处理和分布式训练可以显著加快模型性能并减少开发时间。在下一节中，我们将考虑通过实施优化模型本身的技术来改善 NN 性能的其他方法。
模型优化
模型优化是一个关注 NN 模型的基础实现以及它们如何训练的高级主题。随着这一领域的研究不断发展，PyTorch 已经为模型优化添加了各种功能。在本节中，我们将探讨三个优化领域——超参数调整、量化和剪枝，并为您提供参考代码，供您在自己的设计中使用。
超参数调整
深度学习模型开发通常涉及选择许多用于设计模型和训练模型的变量。这些变量称为超参数，可以包括架构变体，如层数、层深度和核大小，以及可选阶段，如池化或批量归一化。超参数还可能包括损失函数或优化参数的变体，例如 LR 或权重衰减率。
在这一部分，我将向您展示如何使用一个名为 Ray Tune 的包来管理您的超参数优化。研究人员通常会手动测试一小组超参数。然而，Ray Tune 允许您配置您的超参数，并确定哪些设置对性能最佳。
Ray Tune 支持最先进的超参数搜索算法和分布式训练。它不断更新新功能。让我们看看如何使用 Ray Tune 进行超参数调整。
还记得我们在第三章中为图像分类训练的 LeNet5 模型吗？让我们尝试不同的模型配置和训练参数，看看我们是否可以使用超参数调整来改进我们的模型。
为了使用 Ray Tune，我们需要对我们的模型进行以下更改：
-
定义我们的超参数及其搜索空间。
-
编写一个函数来封装我们的训练循环。
-
运行 Ray Tune 超参数调整。
让我们重新定义我们的模型，以便我们可以配置全连接层中节点的数量，如下面的代码所示：
import torch.nn as nn
import torch.nn.functional as F
class Net(nn.Module):
def __init__(self, nodes_1=120, nodes_2=84):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, nodes_1) # ①
self.fc2 = nn.Linear(nodes_1, nodes_2) # ②
self.fc3 = nn.Linear(nodes_2, 10)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 5 * 5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
①
配置fc1中的节点。
②
配置fc2中的节点。
到目前为止，我们有两个超参数，nodes_1和nodes_2。让我们还定义另外两个超参数，lr和batch_size，这样我们就可以在训练中改变学习率和批量大小。
在下面的代码中，我们导入ray包并定义超参数配置：
from ray import tune
import numpy as np
config = {
"nodes_1": tune.sample_from(
lambda _: 2 ** np.random.randint(2, 9)),
"nodes_2": tune.sample_from(
lambda _: 2 ** np.random.randint(2, 9)),
"lr": tune.loguniform(1e-4, 1e-1),
"batch_size": tune.choice([2, 4, 8, 16])
}
在每次运行期间，这些参数的值是从指定的搜索空间中选择的。您可以使用方法tune.sample_from()和一个lambda函数来定义搜索空间，或者您可以使用内置的采样函数。在这种情况下，layer_1和layer_2分别使用sample_from()从2到9中随机选择一个值。
lr和batch_size使用内置函数，其中lr被随机选择为从 1e-4 到 1e-1 的双精度数，batch_size被随机选择为2、4、8或16中的一个。
接下来，我们需要将我们的训练循环封装到一个函数中，该函数以配置字典作为输入。这个训练循环函数将被 Ray Tune 调用。
在编写我们的训练循环之前，让我们定义一个函数来加载 CIFAR-10 数据，这样我们可以在训练期间重复使用来自同一目录的数据。下面的代码类似于我们在第三章中使用的数据加载代码：
import torch
import torchvision
from torchvision import transforms
def load_data(data_dir="./data"):
train_transforms = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(
(0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))])
test_transforms = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(
(0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))])
trainset = torchvision.datasets.CIFAR10(
root=data_dir, train=True,
download=True, transform=train_transforms)
testset = torchvision.datasets.CIFAR10(
root=data_dir, train=False,
download=True, transform=test_transforms)
return trainset, testset
现在我们可以将训练循环封装成一个函数，train_model()，如下面的代码所示。这是一个大段的代码；但是，这应该对您来说很熟悉：
from torch import optim
from torch import nn
from torch.utils.data import random_split
def train_model(config):
device = torch.device("cuda" if
torch.cuda.is_available() else "cpu")
model = Net(config['nodes_1'],
config['nodes_2']).to(device=device) # ①
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=config['lr'],
momentum=0.9) # ②
trainset, testset = load_data()
test_abs = int(len(trainset) * 0.8)
train_subset, val_subset = random_split(
trainset,
[test_abs, len(trainset) - test_abs])
trainloader = torch.utils.data.DataLoader(
train_subset,
batch_size=int(config["batch_size"]),
shuffle=True) # ③
valloader = torch.utils.data.DataLoader(
val_subset,
batch_size=int(config["batch_size"]),
shuffle=True) # ③
for epoch in range(10):
train_loss = 0.0
epoch_steps = 0
for data in trainloader:
inputs, labels = data
inputs = inputs.to(device)
labels = labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
train_loss += loss.item()
val_loss = 0.0
total = 0
correct = 0
for data in valloader:
with torch.no_grad():
inputs, labels = data
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
_, predicted = torch.max(
outputs.data, 1)
total += labels.size(0)
correct += \
(predicted == labels).sum().item()
loss = criterion(outputs, labels)
val_loss += loss.cpu().numpy()
print(f'epoch: {epoch} ',
f'train_loss: ',
f'{train_loss/len(trainloader)}',
f'val_loss: ',
f'{val_loss/len(valloader)}',
f'val_acc: {correct/total}')
tune.report(loss=(val_loss / len(valloader)),
accuracy=correct / total)
①
使模型层可配置。
②
使学习率可配置。
③
使批量大小可配置。
接下来我们想要运行 Ray Tune，但首先我们需要确定我们想要使用的调度程序和报告程序。调度程序确定 Ray Tune 如何搜索和选择超参数，而报告程序指定我们希望如何查看结果。让我们在下面的代码中设置它们：
from ray.tune import CLIReporter
from ray.tune.schedulers import ASHAScheduler
scheduler = ASHAScheduler(
metric="loss",
mode="min",
max_t=10,
grace_period=1,
reduction_factor=2)
reporter = CLIReporter(
metric_columns=["loss",
"accuracy",
"training_iteration"])
对于调度器，我们将使用异步连续减半算法（ASHA）进行超参数搜索，并指示它最小化损失。对于报告器，我们将配置一个 CLI 报告器，以便在每次运行时在 CLI 上报告损失、准确性、训练迭代和选择的超参数。
最后，我们可以使用以下代码中显示的run()方法运行 Ray Tune：
from functools import partial
result = tune.run(
partial(train_model),
resources_per_trial={"cpu": 2, "gpu": 1},
config=config,
num_samples=10,
scheduler=scheduler,
progress_reporter=reporter)
我们提供资源并指定配置。我们传入我们的配置字典，指定样本或运行的数量，并传入我们的scheduler和reporter函数。
Ray Tune 将报告结果。get_best_trial()方法返回一个包含有关最佳试验信息的对象。我们可以打印出产生最佳结果的超参数设置，如下面的代码所示：
best_trial = result.get_best_trial(
"loss", "min", "last")
print("Best trial config: {}".format(
best_trial.config))
print("Best trial final validation loss:",
"{}".format(
best_trial.last_result["loss"]))
print("Best trial final validation accuracy:",
"{}".format(
best_trial.last_result["accuracy"]))
您可能会发现 Ray Tune API 的其他功能有用。表 6-1 列出了tune.schedulers中可用的调度器。
表 6-1. Ray Tune 调度器
| 调度方法 | 描述 |
|---|---|
| ASHA | 运行异步连续减半算法的调度器 |
| HyperBand | 运行 HyperBand 早停算法的调度器 |
| 中位数停止规则 | 基于中位数停止规则的调度器，如“Google Vizier: A Service for Black-Box Optimization”中所述。 |
| 基于人口的训练 | 基于人口训练算法的调度器 |
| 基于人口的训练重放 | 重放人口训练运行的调度器 |
| BOHB | 使用贝叶斯优化和 HyperBand 的调度器 |
| FIFOScheduler | 简单的调度器，按提交顺序运行试验 |
| TrialScheduler | 基于试验的调度器 |
| Shim 实例化 | 基于提供的字符串的调度器 |
更多信息可以在Ray Tune 文档中找到。正如您所看到的，Ray Tune 具有丰富的功能集，但也有其他支持 PyTorch 的超参数包。这些包括Allegro Trains和Optuna。
通过找到最佳设置，超参数调整可以显著提高 NN 模型的性能。接下来，我们将探讨另一种优化模型的技术：量化。
量化
NNs 实现为计算图，它们的计算通常使用 32 位（或在某些情况下，64 位）浮点数。然而，我们可以使我们的计算使用低精度数字，并通过应用量化仍然实现可比较的结果。
量化是指使用低精度数据进行计算和访问内存的技术。这些技术可以减小模型大小，减少内存带宽，并由于内存带宽节省和使用 int8 算术进行更快的推断而执行更快的计算。
一种快速的量化方法是将所有计算精度减半。让我们再次考虑我们的 LeNet5 模型示例，如下面的代码所示：
import torch
from torch import nn
import torch.nn.functional as F
class LeNet5(nn.Module):
def __init__(self):
super(LeNet5, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = F.max_pool2d(
F.relu(self.conv1(x)), (2, 2))
x = F.max_pool2d(
F.relu(self.conv2(x)), 2)
x = x.view(-1,
int(x.nelement() / x.shape[0]))
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
model = LeNet5()
默认情况下，所有计算和内存都实现为 float32。我们可以使用以下代码检查模型参数的数据类型：
for n, p in model.named_parameters():
print(n, ": ", p.dtype)
# out:
# conv1.weight : torch.float32
# conv1.bias : torch.float32
# conv2.weight : torch.float32
# conv2.bias : torch.float32
# fc1.weight : torch.float32
# fc1.bias : torch.float32
# fc2.weight : torch.float32
# fc2.bias : torch.float32
# fc3.weight : torch.float32
# fc3.bias : torch.float32
如预期，我们的数据类型是 float32。然而，我们可以使用half()方法在一行代码中将模型减少到半精度：
model = model.half()
for n, p in model.named_parameters():
print(n, ": ", p.dtype)
# out:
# conv1.weight : torch.float16
# conv1.bias : torch.float16
# conv2.weight : torch.float16
# conv2.bias : torch.float16
# fc1.weight : torch.float16
# fc1.bias : torch.float16
# fc2.weight : torch.float16
# fc2.bias : torch.float16
# fc3.weight : torch.float16
# fc3.bias : torch.float16
现在我们的计算和内存值是 float16。使用half()通常是量化模型的一种快速简单方法。值得一试，看看性能是否适合您的用例。
然而，在许多情况下，我们不希望以相同方式量化每个计算，并且我们可能需要将量化超出 float16 值。对于这些其他情况，PyTorch 提供了三种额外的量化模式：动态量化、训练后静态量化和量化感知训练（QAT）。
当权重的计算或内存带宽限制吞吐量时，使用动态量化。这通常适用于 LSTM、RNN、双向编码器表示来自变压器（BERT）或变压器网络。当激活的内存带宽限制吞吐量时，通常适用于 CNN 的静态量化。当静态量化无法满足精度要求时，使用 QAT。
让我们为每种类型提供一些参考代码。所有类型将权重转换为 int8。它们在处理激活和内存访问方面有所不同。
动态量化是最简单的类型。它会将激活即时转换为 int8。计算使用高效的 int8 值，但激活以浮点格式读取和写入内存。
以下代码向您展示了如何使用动态量化量化模型：
import torch.quantization
quantized_model = \
torch.quantization.quantize_dynamic(
model,
{torch.nn.Linear},
dtype=torch.qint8)
我们所需做的就是传入我们的模型并指定量化层和量化级别。
警告
量化取决于用于运行量化模型的后端。目前，量化运算符仅在以下后端中支持 CPU 推断：x86（fbgemm）和 ARM（qnnpack）。然而，量化感知训练在完全浮点数上进行，并且可以在 GPU 或 CPU 上运行。
后训练静态量化可通过观察训练期间不同激活的分布，并决定在推断时如何量化这些激活来进一步降低延迟。这种类型的量化允许我们在操作之间传递量化值，而无需在内存中来回转换浮点数和整数：
static_quant_model = LeNet5()
static_quant_model.qconfig = \
torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare(
static_quant_model, inplace=True)
torch.quantization.convert(
static_quant_model, inplace=True)
后训练静态量化需要配置和训练以准备使用。我们配置后端以使用 x86（fbgemm），并调用torch.quantization.prepare来插入观察器以校准模型并收集统计信息。然后我们将模型转换为量化版本。
量化感知训练通常会产生最佳精度。在这种情况下，所有权重和激活在训练的前向和后向传递期间都被“伪量化”。浮点值四舍五入为 int8 等效值，但计算仍然以浮点数进行。也就是说，在训练期间进行量化时，权重调整是“知道”的。以下代码显示了如何使用 QAT 量化模型：
qat_model = LeNet5()
qat_mode.qconfig = \
torch.quantization.get_default_qat_qconfig('fbgemm')
torch.quantization.prepare_qat(
qat_model, inplace=True)
torch.quantization.convert(
qat_model, inplace=True)
再次，我们需要配置后端并准备模型，然后调用convert()来量化模型。
PyTorch 的量化功能正在不断发展，目前处于测试阶段。请参考PyTorch 文档获取有关如何使用量化包的最新信息。
修剪
现代深度学习模型可能具有数百万个参数，并且可能难以部署。但是，模型是过度参数化的，参数通常可以减少而几乎不影响准确性或模型性能。修剪是一种通过最小影响性能来减少模型参数数量的技术。这使您可以部署具有更少内存、更低功耗和减少硬件资源的模型。
修剪模型示例
修剪可以应用于nn.module。由于nn.module可能包含单个层、多个层或整个模型，因此可以将修剪应用于单个层、多个层或整个模型本身。让我们考虑我们的 LeNet5 模型示例：
from torch import nn
import torch.nn.functional as F
class LeNet5(nn.Module):
def __init__(self):
super(LeNet5, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = F.max_pool2d(
F.relu(self.conv1(x)), (2, 2))
x = F.max_pool2d(
F.relu(self.conv2(x)), 2)
x = x.view(-1,
int(x.nelement() / x.shape[0]))
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
我们的 LeNet5 模型有五个子模块——conv1、conv2、fc1、fc2和fc3。模型参数包括其权重和偏差，可以使用named_parameters()方法显示。让我们看看conv1层的参数：
device = torch.device("cuda" if
torch.cuda.is_available() else "cpu")
model = LeNet5().to(device)
print(list(model.conv1.named_parameters()))
# out:
# [('weight', Parameter containing:
# tensor([[[[0.0560, 0.0066, ..., 0.0183, 0.0783]]]],
# device='cuda:0',
# requires_grad=True)),
# ('bias', Parameter containing:
# tensor([0.0754, -0.0356, ..., -0.0111, 0.0984],
# device='cuda:0',
# requires_grad=True))]
本地和全局修剪
本地修剪是指我们仅修剪模型的特定部分。通过这种技术，我们可以将本地修剪应用于单个层或模块。只需调用修剪方法，传入层，并设置其选项如下代码所示：
import torch.nn.utils.prune as prune
prune.random_unstructured(model.conv1,
name="weight",
amount=0.25)
此示例将随机非结构化修剪应用于我们模型中conv1层中名为weight的参数。这只修剪了权重参数。我们也可以使用以下代码修剪偏置参数：
prune.random_unstructured(model.conv1,
name="bias",
amount=0.25)
修剪可以进行迭代应用，因此您可以使用不同维度上的其他修剪方法进一步修剪相同的参数。
您可以以不同方式修剪模块和参数。例如，您可能希望按模块或层类型修剪，并将修剪应用于卷积层和线性层的方式不同。以下代码演示了一种方法：
model = LeNet5().to(device)
for name, module in model.named_modules():
if isinstance(module, torch.nn.Conv2d):
prune.random_unstructured(module,
name='weight',
amount=0.3) # ①
elif isinstance(module, torch.nn.Linear):
prune.random_unstructured(module,
name='weight',
amount=0.5) # ②
①
通过 30%修剪所有 2D 卷积层。
②
通过 50%修剪所有线性层。
另一个使用修剪 API 的方法是应用全局修剪，即我们将修剪方法应用于整个模型。例如，我们可以全局修剪我们模型参数的 25%，这可能会导致每个层的不同修剪率。以下代码演示了一种应用全局修剪的方法：
model = LeNet5().to(device)
parameters_to_prune = (
(model.conv1, 'weight'),
(model.conv2, 'weight'),
(model.fc1, 'weight'),
(model.fc2, 'weight'),
(model.fc3, 'weight'),
)
prune.global_unstructured(
parameters_to_prune,
pruning_method=prune.L1Unstructured,
amount=0.25)
在这里我们修剪整个模型中所有参数的 25%。
修剪 API
PyTorch 在其torch.nn.utils.prune模块中提供了对修剪的内置支持。表 6-2 列出了修剪 API 中可用的函数。
表 6-2. 修剪函数
| 函数 | 描述 |
|---|---|
is_pruned(*module*) | 检查模块是否已修剪 |
remove(*module*, *name*) | 从模块中删除修剪重参数化和从前向钩子中删除修剪方法 |
custom_from_mask(*module*, *name*, *mask*) | 通过应用mask中的预先计算的掩码，修剪与module中名为name的参数对应的张量 |
global_unstructured(*params*, *pruning_method*) | 通过应用指定的pruning_method，全局修剪与params中所有参数对应的张量 |
ln_structured(*module*, *name*, *amount*, *n*, *dim*) | 通过移除指定dim上具有最低 Ln-范数的（当前未修剪的）通道，修剪与module中名为name的参数对应的张量中的指定amount |
random_structured(*module*, *name*, *amount*, *dim*) | 通过随机选择指定dim上的通道，移除与module中名为name的参数对应的张量中的指定amount（当前未修剪的）通道 |
l1_unstructured(*module*, *name*, *amount*) | 通过移除具有最低 L1-范数的指定amount（当前未修剪的）单元，修剪与module中名为name的参数对应的张量 |
random_unstructured(*module*, *name*, *amount*) | 通过随机选择指定amount的（当前未修剪的）单元，修剪与module中名为name的参数对应的张量 |
自定义修剪方法
如果找不到适合您需求的修剪方法，您可以创建自己的修剪方法。为此，请从torch.nn.utils.prune中提供的BasePruningMethod类创建一个子类。在大多数情况下，您可以将call()、apply_mask()、apply()、prune()和remove()方法保持原样。
但是，您需要编写自己的__init__()构造函数和compute_mask()方法来描述您的修剪方法如何计算掩码。此外，您需要指定修剪类型（structured、unstructured或global）。以下代码显示了一个示例：
class MyPruningMethod(prune.BasePruningMethod):
PRUNING_TYPE = 'unstructured'
def compute_mask(self, t, default_mask):
mask = default_mask.clone()
mask.view(-1)[::2] = 0
return mask
def my_unstructured(module, name):
MyPruningMethod.apply(module, name)
return module
首先我们定义类。此示例根据compute_mask()中的代码定义了每隔一个参数进行修剪。PRUNING_TYPE用于配置修剪类型为unstructured。然后我们包含并应用一个实例化该方法的函数。您可以按以下方式将此修剪应用于您的模型：
model = LeNet5().to(device)
my_unstructured(model.fc1, name='bias')
您现在已经创建了自己的自定义修剪方法，并可以在本地或全局应用它。
本章向您展示了如何使用 PyTorch 加速培训并优化模型。下一步是将您的模型和创新部署到世界上。在下一章中，您将学习如何将您的模型部署到云端、移动设备和边缘设备，并且我将提供一些参考代码来构建快速应用程序，展示您的设计。
