《动手学深度学习》(五) -- 多GPU训练

1 自动并行

        深度学习框架（例如，MxNet和PyTorch）会在后端自动构建计算图。利用计算图，系统可以了解所有依赖关系，并且可以选择性地并行执行多个不相互依赖的任务以提高速度。

        通常情况下单个操作符将使用所有CPU或单个GPU上的所有计算资源。例如，即使在一台机器上有多个CPU处理器，dot 操作符也将使用所有CPU上的所有核心（和线程）。这样的行为同样适用于单个GPU。因此，并行化对于单设备计算机来说并不是很有用，而并行化对于多个设备就很重要了。虽然并行化通常应用在多个GPU之间，但增加本地CPU以后还将提高少许性能。借助自动并行化框架的便利性，我们可以依靠几行Python代码实现相同的目标。更广泛地考虑，我们对自动并行计算的讨论主要集中在使用CPU和GPU的并行计算上，以及计算和通信的并行化内容。

1.1 基于GPU的并行计算

        在实验开始前，先定义一个度量时间的类，用来衡量（和改进）模型性能时将⾮常有⽤。

class Benchmark:
    """⽤于测量运⾏时间"""
    def __init__(self, description='Done'):
        self.description = description
    
    def __enter__(self):
        self.timer = d2l.Timer()
        return self
    def __exit__(self, *args):
        print(f'{self.description}: {self.timer.stop():.4f} sec')

        接下里，让我们从定义一个具有参考性的用于测试的工作负载开始：下面的run函数将执行$10$ 次“矩阵－矩阵”乘法时需要使用的数据分配到两个变量（x_gpu1和x_gpu2）中，这两个变量分别位于我们选择的不同设备上。

devices = d2l.try_all_gpus()
def run(x):
    return [x.mm(x) for _ in range(50)]

x_gpu1 = torch.rand(size=(4000, 4000), device=devices[0])
x_gpu2 = torch.rand(size=(4000, 4000), device=devices[1])

        现在我们使用函数来处理数据。在测量之前预热设备（对设备执行一次传递）来确保缓存的作用不影响最终的结果。torch.cuda.synchronize()函数将会等待一个CUDA设备上的所有流中的所有核心的计算完成。函数接受一个device参数，代表是哪个设备需要同步。如果device参数是None（默认值），它将使用current_device()找出的当前设备。

run(x_gpu1)
run(x_gpu2)  # 预热设备
torch.cuda.synchronize(devices[0])
torch.cuda.synchronize(devices[1])

with d2l.Benchmark('GPU1 time'):
    run(x_gpu1)
    torch.cuda.synchronize(devices[0])

with d2l.Benchmark('GPU2 time'):
    run(x_gpu2)
    torch.cuda.synchronize(devices[1])

# GPU1 time: 0.5620 sec
# GPU2 time: 0.5732 sec

        如果我们删除两个任务之间的synchronize语句，系统就可以在两个设备上自动实现并行计算。

with d2l.Benchmark('GPU1 & GPU2'):
    run(x_gpu1)
    run(x_gpu2)
    torch.cuda.synchronize()        # GPU1 & GPU2: 0.5573 sec

        在上述情况下，总执行时间小于两个部分执行时间的总和，因为深度学习框架自动调度两个GPU设备上的计算，而不需要用户编写复杂的代码。

1.2 并行计算与通信

        在许多情况下，我们需要在不同的设备之间移动数据，比如在CPU和GPU之间，或者在不同的GPU之间。例如，当我们打算执行分布式优化时，就需要移动数据来聚合多个加速卡上的梯度。通过在GPU上计算，然后将结果复制回CPU来模拟这个过程。

def copy_to_cpu(x, non_blocking=False):
    return [y.to('cpu', non_blocking=non_blocking) for y in x]

with d2l.Benchmark('在GPU1上运行'):
    y = run(x_gpu1)
    torch.cuda.synchronize()

with d2l.Benchmark('复制到CPU'):
    y_cpu = copy_to_cpu(y)
    torch.cuda.synchronize()
# 在GPU1上运行: 0.5355 sec
# 复制到CPU: 2.1120 sec

        这种方式效率不高。当列表中的其余部分还在计算时，我们可能就已经开始将y的部分复制到CPU了。例如，当我们计算一个小批量的（反传）梯度时。某些参数的梯度将比其他参数的梯度更早可用。因此，在GPU仍在运行时就开始使用PCI-Express总线带宽来移动数据对我们是有利的。在PyTorch中，to()和copy_()等函数都允许显式的non_blocking参数，这允许在不需要同步时调用方可以绕过同步。设置non_blocking=True让我们模拟这个场景。

with d2l.Benchmark('在GPU1上运行并复制到CPU'):
    y = run(x_gpu1)
    y_cpu = copy_to_cpu(y, True)
    torch.cuda.synchronize()        # 在GPU1上运行并复制到CPU: 1.5400 sec

        两个操作所需的总时间少于它们各部分操作所需时间的总和。请注意，与并行计算的区别是通信操作使用的资源：CPU和GPU之间的总线。事实上，我们可以在两个设备上同时进行计算和通信。如上所述，计算和通信之间存在的依赖关系是必须先计算y[i]，然后才能将其复制到CPU。幸运的是，系统可以在计算y[i]的同时复制y[i-1]，以减少总的运行时间。

---

2 多GPU训练

2.1 数据并行性

        假设一台机器有$k$个GPU。给定需要训练的模型，虽然每个GPU上的参数值都是相同且同步的，但是每个GPU都将独立地维护一组完整的模型参数。例如， 下图演示了在$k=2$时基于数据并行方法训练模型。

利⽤两个GPU上的数据，并⾏计算小批量随机梯度下降

        一般来说，$k$个GPU并行训练过程如下：

• 在任何一次训练迭代中，给定的随机的小批量样本都将被分成$k$个部分，并均匀地分配到GPU上。
• 每个GPU根据分配给它的小批量子集，计算模型参数的损失和梯度。
• 将$k$个GPU中的局部梯度聚合，以获得当前小批量的随机梯度。
• 聚合梯度被重新分发到每个GPU中。
• 每个GPU使用这个小批量随机梯度，来更新它所维护的完整的模型参数集。

        在实践中请注意，当在$k$个GPU上训练时，需要扩大小批量的大小为$k$的倍数，这样每个GPU都有相同的工作量，就像只在单个GPU上训练一样。因此，在16-GPU服务器上可以显著地增加小批量数据量的大小，同时可能还需要相应地提高学习率。下面我们将使用一个简单网络来演示多GPU训练。

import torch
from torch import nn
from torch.nn import functional as F 
from d2l import torch as d2l

2.2 简单网络

        为了详细说明参数交换和同步，这里对LeNet稍加修改，并从零开始定义。

# 初始化模型参数
scale = 0.01
w1 = torch.randn(size=(20, 1, 3, 3)) * scale
b1 = torch.zeros(20)
w2 = torch.randn(size=(50, 20, 5, 5)) * scale
b2 = torch.zeros(50)
w3 = torch.randn(size=(800, 128)) * scale
b3 = torch.zeros(128)
w4 = torch.randn(size=(128, 10)) * scale
b4 = torch.zeros(10)
params = [w1, b1, w2, b2, w3, b3, w4, b4]

# 定义模型
def lenet(X, params):
    h1_conv = F.conv2d(input=X, weight=params[0], bias=params[1])
    h1_activation = F.relu(h1_conv)
    h1 = F.avg_pool2d(input=h1_activation, kernel_size=(2, 2), stride=(2, 2))
    h2_conv = F.conv2d(input=h1, weight=params[2], bias=params[3])
    h2_activation = F.relu(h2_conv)
    h2 = F.avg_pool2d(input=h2_activation, kernel_size=(2, 2), stride=(2, 2))
    h2 = h2.reshape(h2.shape[0], -1)
    h3_linear = torch.mm(h2, params[4]) + params[5]
    h3 = F.relu(h3_linear)
    y_hat = torch.mm(h3, params[6]) + params[7]
    return y_hat

# 交叉熵函数
loss = nn.CrossEntropyLoss(reduction='none')

        补充：简单易懂的利用F.conv2d函数进行卷积的前向传播和反向传播

torch.nn.functional.conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1)

input – input tensor of shape (minibatch, in_channels, iH, iW)
weight – filters of shape(out_channels, $\frac{in\_channels}{groups}$, kH, kW)
output – output tensor of shape (minibatch, out_channels, oH, oW)

2.3 数据同步

        对于高效的多GPU训练，我们需要两个基本操作。首先，我们需要[向多个设备分发参数]并附加梯度（get_params）。如果没有参数，就不可能在GPU上评估网络。其次，需要跨多个设备对参数求和，也就是说，需要一个allreduce函数。

def get_params(params, device):
    new_params = [p.to(device) for p in params]
    for p in new_params:
        p.requires_grad_()
    return new_params

        通过将模型参数复制到一个GPU。

new_params = get_params(params, d2l.try_gpu(0))
print('b1 权重：', new_params[1])
print('b1 梯度：', new_params[1].grad)

b1 权重： tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], device=‘cuda:0’, requires_grad=True)
b1 梯度： None

        由于还没有进行任何计算，因此权重参数的梯度仍然为零。假设现在有一个向量分布在多个GPU上，下面的[allreduce函数将所有向量相加，并将结果广播给所有GPU]。请注意，我们需要将数据复制到累积结果的设备，才能使函数正常工作。

def allreduce(data):
    for i in range(1, len(data)):       # 先在GPU0上求和
        data[0][:] += data[i].to(data[0].device)
    for i in range(1, len(data)):       # 广播给所有GPU
        data[i][:] = data[0].to(data[i].device)

        通过在不同设备上创建具有不同值的向量并聚合它们。

data = [torch.ones((1, 2), device=d2l.try_gpu(i)) * (i + 1) for i in range(2)]
print('allreduce之前：\n', data[0], '\n', data[1])
allreduce(data)
print('allreduce之后：\n', data[0], '\n', data[1])

2.4 数据分发

        我们需要一个简单的工具函数，[将一个小批量数据均匀地分布在多个GPU上]。例如，有两个GPU时，我们希望每个GPU可以复制一半的数据。因为深度学习框架的内置函数编写代码更方便、更简洁，所以在$4\times 5$矩阵上使用它进行尝试。

data = torch.arange(20).reshape(4, 5)
devices = [torch.device('cuda:0'), torch.device('cuda:1')]
split = nn.parallel.scatter(data, devices)			# 将输入数据若干等分
print('input:', data)
print('load into', devices)
print('output:', split)

        为了方便以后复用，我们定义了可以同时拆分数据和标签的split_batch函数。

def split_batch(X, y, devices):
    """将X和y拆分到多个设备上"""
    assert X.shape[0] == y.shape[0]
    return (nn.parallel.scatter(X, devices), nn.parallel.scatter(y, devices))

2.5 训练

        现在我们可以[在一个小批量上实现多GPU训练]。在多个GPU之间同步数据将使用刚才讨论的辅助函数allreduce和split_and_load。我们不需要编写任何特定的代码来实现并行性。因为计算图在小批量内的设备之间没有任何依赖关系，因此它是“自动地”并行执行。

def train_batch(X, y, device_params, devices, lr):
    X_shards, y_shards = split_batch(X, y, devices)
    # 在每个GPU上分别计算损失
    ls = [loss(lenet(X_shard, device_w), y_shard).sum()
          for X_shard, y_shard, device_w in zip(X_shards, y_shards, device_params)]
    for l in ls:    # 反向传播在每个GPU上分别执行
        l.backward()
    # 将每个GPU的所有梯度相加，并将其广播到所有GPU
    with torch.no_grad():
        for i in range(len(device_params[0])):
            allreduce([device_params[c][i].grad for c in range(len(devices))])

    # 在每个GPU上分别更新模型参数
    for param in device_params:
        d2l.sgd(param, lr, X.shape[0])  # 使用全尺寸的小批量

        现在，我们可以[定义训练函数]。与前几章中略有不同：训练函数需要分配GPU并将所有模型参数复制到所有设备。显然，每个小批量都是使用train_batch函数来处理多个GPU。我们只在一个GPU上计算模型的精确度，而让其他GPU保持空闲，尽管这是相对低效的，但是使用方便且代码简洁。

from torchvision import transforms
from torch.utils.data import DataLoader
def train(num_gpus, batch_size, lr):
    train_loader, test_loader = d2l.load_data_fashion_mnist(batch_size)
    devices = [d2l.try_gpu(i) for i in range(num_gpus)]
    device_params = [get_params(params, d) for d in devices]
    num_epochs = 10
    animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
    timer = d2l.Timer()
    for epoch in range(num_epochs):
        timer.start()
        for X, y in train_loader:
            # 为单个小批量执行多GPU训练
            train_batch(X, y, device_params, devices, lr)
            torch.cuda.synchronize()
        timer.stop()
        # 在GPU0上评估模型
        animator.add(epoch+1, (d2l.evaluate_accuracy_gpu(lambda x: lenet(x, device_params[0]), test_loader, devices[0])))
    print(f'测试精度：{animator.Y[0][-1]:.2f}, {timer.avg():.1f}秒/轮, 在{str(devices)}')

        让我们看看[在单个GPU上运行]效果得有多好。首先使用的批量大小是$256$，学习率是$0.2$。

train(num_gpus=1, batch_size=256, lr=0.2)   # 测试精度：0.82, 2.5秒/轮, 在[device(type='cuda', index=0)]

        保持批量大小和学习率不变，并[增加为2个GPU]，我们可以看到测试精度与之前的实验基本相同。不同的GPU个数在算法寻优方面是相同的。不幸的是，这里没有任何有意义的加速：模型实在太小了；而且数据集也太小了，在这个数据集中，我们实现的多GPU训练的简单方法受到了巨大的Python开销的影响。在未来，我们将遇到更复杂的模型和更复杂的并行化方法。尽管如此，让我们看看Fashion-MNIST数据集上会发生什么。

train(num_gpus=2, batch_size=256, lr=0.2)
# 测试精度：0.83, 3.4秒/轮, 在[device(type='cuda', index=0), device(type='cuda', index=1)]

---

3 多GPU的简洁实现

        下面将展示如何使用深度学习框架的高级API来实现优化同步工具以获得高性能。

3.1 简单网络

        这里选择的是ResNet-18，由于输入的图像很小，这里做了一些修改，使用了更小的卷积核、步长和填充，删除了最大汇聚层。

def resnet18(num_classes, in_channels=1):
    """稍加修改的ResNet-18模型"""   
    def resnet_block(in_channels, out_channels, num_residuals, first_block=False):
        blk = []
        for i in range(num_residuals):
            if i == 0 and not first_block:
                blk.append(d2l.Residual(in_channels, out_channels, use_1x1conv=True, strides=2))
            else:
                blk.append(d2l.Residual(out_channels, out_channels))
        return nn.Sequential(*blk)

    # 该模型使⽤了更⼩的卷积核、步⻓和填充，⽽且删除了最⼤汇聚层
    net = nn.Sequential(nn.Conv2d(in_channels, 64, kernel_size=3, stride=1, padding=1),
                        nn.BatchNorm2d(64),
                        nn.ReLU())
    net.add_module('resnet_block1', resnet_block(64, 64, 2, first_block=True))
    net.add_module('resnet_block2', resnet_block(64, 128, 2))
    net.add_module('resnet_block3', resnet_block(128, 256, 2))
    net.add_module('resnet_block4', resnet_block(256, 512, 2))
    net.add_module('global_avg_pool', nn.AdaptiveAvgPool2d((1, 1)))
    net.add_module('fc', nn.Sequential(nn.Flatten(), 
                                    nn.Linear(512, num_classes)))
    return net

3.2 网络初始化

        在训练回路中初始化⽹络。

net = resnet18(10)
# 获取GPU列表
devices = d2l.try_all_gpus()
# 我们将在训练代码实现中初始化网络

3.3 训练

        用于训练的代码需要执行几个基本功能才能实现高效并行：

• 需要在所有设备上初始化网络参数。
• 在数据集上迭代时，要将小批量数据分配到所有设备上。
• 跨设备并行计算损失及其梯度。
• 聚合梯度，并相应地更新参数。

        最后，并行地计算精确度和发布网络的最终性能。除了需要拆分和聚合数据外，训练代码与前几章的实现非常相似。

def train(net, num_gpus, batch_size, lr):
    train_loader, test_loader = d2l.load_data_fashion_mnist(batch_size)
    devices = [d2l.try_gpu(i) for i in range(num_gpus)]
    def init_weights(m):
        if type(m) in [nn.Linear, nn.Conv2d]:
            nn.init.normal_(m.weight, std=0.01)
    net.apply(init_weights)
    # 在多个GPU上设置模型
    net = nn.DataParallel(net, device_ids=devices)
    optimizer = torch.optim.SGD(net.parameters(), lr)
    criterion = nn.CrossEntropyLoss()
    timer, num_epochs = d2l.Timer(), 10
    animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
    for epoch in range(num_epochs):
        net.train()
        timer.start()
        for X, y in train_loader:
            optimizer.zero_grad()
            X, y = X.to(devices[0]), y.to(devices[0])
            loss = criterion(net(X), y)
            loss.backward()
            optimizer.step()
        timer.stop()
        animator.add(epoch+1, (d2l.evaluate_accuracy_gpu(net, test_loader), ))
    print(f'测试精度：{animator.Y[0][-1]:.2f}, {timer.avg():.1f}秒/轮, 在{str(devices)}')

        先[在单个GPU上训练网络]进行预热。

train(net, num_gpus=1, batch_size=256, lr=0.1)      # 测试精度：0.91, 18.1秒/轮, 在[device(type='cuda', index=0)]

        接下来我们[使用2个GPU进行训练]。与上一节中评估的LeNet相比，ResNet-18的模型要复杂得多。这就是显示并行化优势的地方，计算所需时间明显大于同步参数需要的时间。因为并行化开销的相关性较小，因此这种操作提高了模型的可伸缩性。

train(net, num_gpus=2, batch_size=512, lr=0.2)      # 测试精度：0.92, 12.6秒/轮, 在[device(type='cuda', index=0), device(type='cuda', index=1)]

        最后，又尝试了[使用4个GPU进行训练]。

train(net, num_gpus=4, batch_size=1024, lr=0.4)
# 测试精度：0.92, 13.0秒/轮, 在[device(type='cuda', index=0), device(type='cuda', index=1), device(type='cuda', index=2), device(type='cuda', index=3)]

        在实际训练中，也要根据模型的复杂度去选择并行GPU的数量，因为数据分发也需要时间消耗，所以要在网络复杂度和数据并行之间综合考虑。