多GPU的简洁实现

:label:sec_multi_gpu_concise

每个新模型的并行计算都从零开始实现是无趣的。此外，优化同步工具以获得高性能也是有好处的。下面我们将展示如何使用深度学习框架的高级API来实现这一点。数学和算法与 :numref:sec_multi_gpu中的相同。不出所料，你至少需要两个GPU来运行本节的代码。

```{.python .input} from d2l import mxnet as d2l from mxnet import autograd, gluon, init, np, npx from mxnet.gluon import nn npx.set_np()

  2. ```{.python .input}
  3. #@tab pytorch
  4. from d2l import torch as d2l
  5. import torch
  6. from torch import nn

[简单网络]

让我们使用一个比 :numref:sec_multi_gpu的LeNet更有意义的网络，它依然能够容易地和快速地训练。我们选择的是 :cite:He.Zhang.Ren.ea.2016中的ResNet-18。因为输入的图像很小，所以稍微修改了一下。与 :numref:sec_resnet的区别在于，我们在开始时使用了更小的卷积核、步长和填充，而且删除了最大汇聚层。

```{.python .input}

@save

def resnet18(num_classes): “””稍加修改的ResNet-18模型””” def resnet_block(num_channels, num_residuals, first_block=False): blk = nn.Sequential() for i in range(num_residuals): if i == 0 and not first_block: blk.add(d2l.Residual( num_channels, use_1x1conv=True, strides=2)) else: blk.add(d2l.Residual(num_channels)) return blk

  1. net = nn.Sequential()
  2. # 该模型使用了更小的卷积核、步长和填充，而且删除了最大汇聚层
  3. net.add(nn.Conv2D(64, kernel_size=3, strides=1, padding=1),
  4.         nn.BatchNorm(), nn.Activation('relu'))
  5. net.add(resnet_block(64, 2, first_block=True),
  6.         resnet_block(128, 2),
  7.         resnet_block(256, 2),
  8.         resnet_block(512, 2))
  9. net.add(nn.GlobalAvgPool2D(), nn.Dense(num_classes))
  10. return net
  2. ```{.python .input}
  3. #@tab pytorch
  4. #@save
  5. def resnet18(num_classes, in_channels=1):
  6.     """稍加修改的ResNet-18模型"""
  7.     def resnet_block(in_channels, out_channels, num_residuals,
  8.                      first_block=False):
  9.         blk = []
  10.         for i in range(num_residuals):
  11.             if i == 0 and not first_block:
  12.                 blk.append(d2l.Residual(in_channels, out_channels,
  13.                                         use_1x1conv=True, strides=2))
  14.             else:
  15.                 blk.append(d2l.Residual(out_channels, out_channels))
  16.         return nn.Sequential(*blk)
  18.     # 该模型使用了更小的卷积核、步长和填充，而且删除了最大汇聚层
  19.     net = nn.Sequential(
  20.         nn.Conv2d(in_channels, 64, kernel_size=3, stride=1, padding=1),
  21.         nn.BatchNorm2d(64),
  22.         nn.ReLU())
  23.     net.add_module("resnet_block1", resnet_block(
  24.         64, 64, 2, first_block=True))
  25.     net.add_module("resnet_block2", resnet_block(64, 128, 2))
  26.     net.add_module("resnet_block3", resnet_block(128, 256, 2))
  27.     net.add_module("resnet_block4", resnet_block(256, 512, 2))
  28.     net.add_module("global_avg_pool", nn.AdaptiveAvgPool2d((1,1)))
  29.     net.add_module("fc", nn.Sequential(nn.Flatten(),
  30.                                        nn.Linear(512, num_classes)))
  31.     return net

网络初始化

:begin_tab:mxnet initialize函数允许我们在所选设备上初始化参数。请参阅 :numref:sec_numerical_stability复习初始化方法。这个函数在多个设备上初始化网络时特别方便。让我们在实践中试一试它的运作方式。 :end_tab:

:begin_tab:pytorch 我们将在训练回路中初始化网络。请参见 :numref:sec_numerical_stability复习初始化方法。 :end_tab:

```{.python .input} net = resnet18(10)

获取GPU列表

devices = d2l.try_all_gpus()

初始化网络的所有参数

net.initialize(init=init.Normal(sigma=0.01), ctx=devices)

  2. ```{.python .input}
  3. #@tab pytorch
  4. net = resnet18(10)
  5. # 获取GPU列表
  6. devices = d2l.try_all_gpus()
  7. # 我们将在训练代码实现中初始化网络

:begin_tab:mxnet 使用 :numref:sec_multi_gpu中引入的split_and_load函数可以切分一个小批量数据，并将切分后的分块数据复制到devices变量提供的设备列表中。网络实例自动使用适当的GPU来计算前向传播的值。我们将在下面生成$4$个观测值，并在GPU上将它们拆分。 :end_tab:

```{.python .input} x = np.random.uniform(size=(4, 1, 28, 28)) x_shards = gluon.utils.split_and_load(x, devices) net(x_shards[0]), net(x_shards[1])

  2. :begin_tab:`mxnet`
  3. 一旦数据通过网络，网络对应的参数就会在*有数据通过的设备上初始化*。这意味着初始化是基于每个设备进行的。由于我们选择的是GPU0和GPU1，所以网络只在这两个GPU上初始化，而不是在CPU上初始化。事实上，CPU上甚至没有这些参数。我们可以通过打印参数和观察可能出现的任何错误来验证这一点。
  4. :end_tab:
  6. ```{.python .input}
  7. weight = net[0].params.get('weight')
  9. try:
  10.     weight.data()
  11. except RuntimeError:
  12.     print('not initialized on cpu')
  13. weight.data(devices[0])[0], weight.data(devices[1])[0]

:begin_tab:mxnet 接下来，让我们使用[在多个设备上并行工作]的代码来替换前面的[评估模型]的代码。 这里主要是 :numref:sec_lenet的evaluate_accuracy_gpu函数的替代，代码的主要区别在于在调用网络之前拆分了一个小批量，其他在本质上是一样的。 :end_tab:

```{.python .input}

@save

def evaluate_accuracy_gpus(net, data_iter, split_f=d2l.split_batch): “””使用多个GPU计算数据集上模型的精度”””

  1. # 查询设备列表
  2. devices = list(net.collect_params().values())[0].list_ctx()
  3. # 正确预测的数量，预测的总数量
  4. metric = d2l.Accumulator(2)
  5. for features, labels in data_iter:
  6.     X_shards, y_shards = split_f(features, labels, devices)
  7.     # 并行运行
  8.     pred_shards = [net(X_shard) for X_shard in X_shards]
  9.     metric.add(sum(float(d2l.accuracy(pred_shard, y_shard)) for
  10.                    pred_shard, y_shard in zip(
  11.                        pred_shards, y_shards)), labels.size)
  12. return metric[0] / metric[1]
  2. ## [**训练**]
  4. 如前所述，用于训练的代码需要执行几个基本功能才能实现高效并行：
  6. * 需要在所有设备上初始化网络参数。
  7. * 在数据集上迭代时，要将小批量数据分配到所有设备上。
  8. * 跨设备并行计算损失及其梯度。
  9. * 聚合梯度，并相应地更新参数。
  11. 最后，并行地计算精确度和发布网络的最终性能。除了需要拆分和聚合数据外，训练代码与前几章的实现非常相似。
  13. ```{.python .input}
  14. def train(num_gpus, batch_size, lr):
  15.     train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
  16.     ctx = [d2l.try_gpu(i) for i in range(num_gpus)]
  17.     net.initialize(init=init.Normal(sigma=0.01), ctx=ctx, force_reinit=True)
  18.     trainer = gluon.Trainer(net.collect_params(), 'sgd',
  19.                             {'learning_rate': lr})
  20.     loss = gluon.loss.SoftmaxCrossEntropyLoss()
  21.     timer, num_epochs = d2l.Timer(), 10
  22.     animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
  23.     for epoch in range(num_epochs):
  24.         timer.start()
  25.         for features, labels in train_iter:
  26.             X_shards, y_shards = d2l.split_batch(features, labels, ctx)
  27.             with autograd.record():
  28.                 ls = [loss(net(X_shard), y_shard) for X_shard, y_shard
  29.                       in zip(X_shards, y_shards)]
  30.             for l in ls:
  31.                 l.backward()
  32.             trainer.step(batch_size)
  33.         npx.waitall()
  34.         timer.stop()
  35.         animator.add(epoch + 1, (evaluate_accuracy_gpus(net, test_iter),))
  36.     print(f'测试精度：{animator.Y[0][-1]:.2f}，{timer.avg():.1f}秒/轮，'
  37.           f'在{str(ctx)}')

```{.python .input}

@tab pytorch

def train(net, numgpus, batch_size, lr): train_iter, test_iter = 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)

  1. # 在多个GPU上设置模型
  2. net = nn.DataParallel(net, device_ids=devices)
  3. trainer = torch.optim.SGD(net.parameters(), lr)
  4. loss = nn.CrossEntropyLoss()
  5. timer, num_epochs = d2l.Timer(), 10
  6. animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
  7. for epoch in range(num_epochs):
  8.     net.train()
  9.     timer.start()
  10.     for X, y in train_iter:
  11.         trainer.zero_grad()
  12.         X, y = X.to(devices[0]), y.to(devices[0])
  13.         l = loss(net(X), y)
  14.         l.backward()
  15.         trainer.step()
  16.     timer.stop()
  17.     animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(net, test_iter),))
  18. print(f'测试精度：{animator.Y[0][-1]:.2f}，{timer.avg():.1f}秒/轮，'
  19.       f'在{str(devices)}')
  2. 让我们看看这在实践中是如何运作的。我们先[**在单个GPU上训练网络**]进行预热。
  4. ```{.python .input}
  5. train(num_gpus=1, batch_size=256, lr=0.1)

```{.python .input}

@tab pytorch

train(net, num_gpus=1, batch_size=256, lr=0.1)

  2. 接下来我们[**使用2个GPU进行训练**]。与 :numref:`sec_multi_gpu`中评估的LeNet相比，ResNet-18的模型要复杂得多。这就是显示并行化优势的地方，计算所需时间明显大于同步参数需要的时间。因为并行化开销的相关性较小，因此这种操作提高了模型的可伸缩性。
  4. ```{.python .input}
  5. train(num_gpus=2, batch_size=512, lr=0.2)

```{.python .input}

@tab pytorch

train(net, num_gpus=2, batch_size=512, lr=0.2) ```

小结

:begin_tab:mxnet

  • Gluon通过提供一个上下文列表，为跨多个设备的模型初始化提供原语。
  • 神经网络可以在（可找到数据的）单GPU上进行自动评估。
  • 每台设备上的网络需要先初始化，然后再尝试访问该设备上的参数，否则会遇到错误。
  • 优化算法在多个GPU上自动聚合。 :end_tab:

:begin_tab:pytorch

  • 神经网络可以在（可找到数据的）单GPU上进行自动评估。
  • 每台设备上的网络需要先初始化，然后再尝试访问该设备上的参数，否则会遇到错误。
  • 优化算法在多个GPU上自动聚合。 :end_tab:

练习

:begin_tab:mxnet

  1. 本节使用ResNet-18，请尝试不同的迭代周期数、批量大小和学习率，以及使用更多的GPU进行计算。如果使用$16$个GPU（例如，在AWS p2.16xlarge实例上）尝试此操作，会发生什么？
  2. 有时候不同的设备提供了不同的计算能力，我们可以同时使用GPU和CPU，那应该如何分配工作？为什么？
  3. 如果去掉npx.waitall()会怎样？你将如何修改训练，以使并行操作最多有两个步骤重叠？ :end_tab:

:begin_tab:pytorch

  1. 本节使用ResNet-18，请尝试不同的迭代周期数、批量大小和学习率，以及使用更多的GPU进行计算。如果使用$16$个GPU（例如，在AWS p2.16xlarge实例上）尝试此操作，会发生什么？
  2. 有时候不同的设备提供了不同的计算能力，我们可以同时使用GPU和CPU，那应该如何分配工作？为什么？ :end_tab:

:begin_tab:mxnet Discussions :end_tab:

:begin_tab:pytorch Discussions :end_tab: