12.6. 多GPU的简洁实现¶ Open the notebook in SageMaker Studio Lab
每个新模型的并行计算都从零开始实现是无趣的。此外，优化同步工具以获得高性能也是有好处的。下面我们将展示如何使用深度学习框架的高级API来实现这一点。数学和算法与 12.5节中的相同。本节的代码至少需要两个GPU来运行。
from mxnet import autograd, gluon, init, np, npx
from mxnet.gluon import nn
from d2l import mxnet as d2l
npx.set_np()
import torch
from torch import nn
from d2l import torch as d2l
import warnings
from d2l import paddle as d2l
warnings.filterwarnings("ignore")
import paddle
from paddle import nn
12.6.1. 简单网络¶
让我们使用一个比 12.5节的LeNet更有意义的网络，它依然能够容易地和快速地训练。我们选择的是 (He et al., 2016)中的ResNet-18。因为输入的图像很小，所以稍微修改了一下。与 7.6节的区别在于，我们在开始时使用了更小的卷积核、步长和填充，而且删除了最大汇聚层。
#@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
net = nn.Sequential()
# 该模型使用了更小的卷积核、步长和填充，而且删除了最大汇聚层
net.add(nn.Conv2D(64, kernel_size=3, strides=1, padding=1),
nn.BatchNorm(), nn.Activation('relu'))
net.add(resnet_block(64, 2, first_block=True),
resnet_block(128, 2),
resnet_block(256, 2),
resnet_block(512, 2))
net.add(nn.GlobalAvgPool2D(), nn.Dense(num_classes))
return net
#@save
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
#@save
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_sublayer("resnet_block1", resnet_block(
64, 64, 2, first_block=True))
net.add_sublayer("resnet_block2", resnet_block(64, 128, 2))
net.add_sublayer("resnet_block3", resnet_block(128, 256, 2))
net.add_sublayer("resnet_block4", resnet_block(256, 512, 2))
net.add_sublayer("global_avg_pool", nn.AdaptiveAvgPool2D((1, 1)))
net.add_sublayer("fc", nn.Sequential(nn.Flatten(),
nn.Linear(512, num_classes)))
return net
12.6.2. 网络初始化¶
initialize函数允许我们在所选设备上初始化参数。请参阅
4.8节复习初始化方法。这个函数在多个设备上初始化网络时特别方便。下面在实践中试一试它的运作方式。
net = resnet18(10)
# 获取GPU列表
devices = d2l.try_all_gpus()
# 初始化网络的所有参数
net.initialize(init=init.Normal(sigma=0.01), ctx=devices)
[07:14:36] ../src/storage/storage.cc:196: Using Pooled (Naive) StorageManager for CPU
[07:14:36] ../src/storage/storage.cc:196: Using Pooled (Naive) StorageManager for GPU
[07:14:36] ../src/storage/storage.cc:196: Using Pooled (Naive) StorageManager for GPU
使用
12.5节中引入的split_and_load函数可以切分一个小批量数据，并将切分后的分块数据复制到devices变量提供的设备列表中。网络实例自动使用适当的GPU来计算前向传播的值。我们将在下面生成\(4\)个观测值，并在GPU上将它们拆分。
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])
[07:14:37] ../src/operator/cudnn_ops.cc:318: Auto-tuning cuDNN op, set MXNET_CUDNN_AUTOTUNE_DEFAULT to 0 to disable
[07:14:37] ../src/operator/cudnn_ops.cc:318: Auto-tuning cuDNN op, set MXNET_CUDNN_AUTOTUNE_DEFAULT to 0 to disable
[07:14:37] ../src/operator/cudnn_ops.cc:318: Auto-tuning cuDNN op, set MXNET_CUDNN_AUTOTUNE_DEFAULT to 0 to disable
[07:14:38] ../src/operator/cudnn_ops.cc:318: Auto-tuning cuDNN op, set MXNET_CUDNN_AUTOTUNE_DEFAULT to 0 to disable
(array([[ 2.2610207e-06, 2.2045981e-06, -5.4046786e-06, 1.2869955e-06,
5.1373163e-06, -3.8297967e-06, 1.4339059e-07, 5.4683451e-06,
-2.8279192e-06, -3.9651104e-06],
[ 2.0698672e-06, 2.0084667e-06, -5.6382510e-06, 1.0498458e-06,
5.5506434e-06, -4.1065491e-06, 6.0830087e-07, 5.4521784e-06,
-3.7365021e-06, -4.1891640e-06]], ctx=gpu(0)),
array([[ 2.4629783e-06, 2.6015525e-06, -5.4362617e-06, 1.2938218e-06,
5.6387889e-06, -4.1360108e-06, 3.5758853e-07, 5.5125256e-06,
-3.1957325e-06, -4.2976326e-06],
[ 1.9431673e-06, 2.2600434e-06, -5.2698201e-06, 1.4807417e-06,
5.4830934e-06, -3.9678889e-06, 7.5751018e-08, 5.6764356e-06,
-3.2530229e-06, -4.0943951e-06]], ctx=gpu(1)))
一旦数据通过网络，网络对应的参数就会在有数据通过的设备上初始化。这意味着初始化是基于每个设备进行的。由于我们选择的是GPU0和GPU1，所以网络只在这两个GPU上初始化，而不是在CPU上初始化。事实上，CPU上甚至没有这些参数。我们可以通过打印参数和观察可能出现的任何错误来验证这一点。
weight = net[0].params.get('weight')
try:
weight.data()
except RuntimeError:
print('not initialized on cpu')
weight.data(devices[0])[0], weight.data(devices[1])[0]
not initialized on cpu
(array([[[ 0.01382882, -0.01183044, 0.01417865],
[-0.00319718, 0.00439528, 0.02562625],
[-0.00835081, 0.01387452, -0.01035946]]], ctx=gpu(0)),
array([[[ 0.01382882, -0.01183044, 0.01417865],
[-0.00319718, 0.00439528, 0.02562625],
[-0.00835081, 0.01387452, -0.01035946]]], ctx=gpu(1)))
接下来，让我们使用在多个设备上并行工作的代码来替换前面的评估模型的代码。
这里主要是
6.6节的evaluate_accuracy_gpu函数的替代，代码的主要区别在于在调用网络之前拆分了一个小批量，其他在本质上是一样的。
#@save
def evaluate_accuracy_gpus(net, data_iter, split_f=d2l.split_batch):
"""使用多个GPU计算数据集上模型的精度"""
# 查询设备列表
devices = list(net.collect_params().values())[0].list_ctx()
# 正确预测的数量，预测的总数量
metric = d2l.Accumulator(2)
for features, labels in data_iter:
X_shards, y_shards = split_f(features, labels, devices)
# 并行运行
pred_shards = [net(X_shard) for X_shard in X_shards]
metric.add(sum(float(d2l.accuracy(pred_shard, y_shard)) for
pred_shard, y_shard in zip(
pred_shards, y_shards)), labels.size)
return metric[0] / metric[1]
我们将在训练回路中初始化网络。请参见 4.8节复习初始化方法。
net = resnet18(10)
# 获取GPU列表
devices = d2l.try_all_gpus()
# 我们将在训练代码实现中初始化网络
net = resnet18(10)
# 获取GPU列表
devices = d2l.try_all_gpus()
# 我们将在训练代码实现中初始化网络
W0818 09:28:17.822042 95393 gpu_resources.cc:61] Please NOTE: device: 0, GPU Compute Capability: 7.0, Driver API Version: 11.8, Runtime API Version: 11.8
W0818 09:28:17.852774 95393 gpu_resources.cc:91] device: 0, cuDNN Version: 8.7.
12.6.3. 训练¶
如前所述，用于训练的代码需要执行几个基本功能才能实现高效并行：
需要在所有设备上初始化网络参数；
在数据集上迭代时，要将小批量数据分配到所有设备上；
跨设备并行计算损失及其梯度；
聚合梯度，并相应地更新参数。
最后，并行地计算精确度和发布网络的最终性能。除了需要拆分和聚合数据外，训练代码与前几章的实现非常相似。
def train(num_gpus, batch_size, lr):
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
ctx = [d2l.try_gpu(i) for i in range(num_gpus)]
net.initialize(init=init.Normal(sigma=0.01), ctx=ctx, force_reinit=True)
trainer = gluon.Trainer(net.collect_params(), 'sgd',
{'learning_rate': lr})
loss = gluon.loss.SoftmaxCrossEntropyLoss()
timer, num_epochs = d2l.Timer(), 10
animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
for epoch in range(num_epochs):
timer.start()
for features, labels in train_iter:
X_shards, y_shards = d2l.split_batch(features, labels, ctx)
with autograd.record():
ls = [loss(net(X_shard), y_shard) for X_shard, y_shard
in zip(X_shards, y_shards)]
for l in ls:
l.backward()
trainer.step(batch_size)
npx.waitall()
timer.stop()
animator.add(epoch + 1, (evaluate_accuracy_gpus(net, test_iter),))
print(f'测试精度：{animator.Y[0][-1]:.2f}，{timer.avg():.1f}秒/轮，'
f'在{str(ctx)}')
def train(net, num_gpus, 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)
# 在多个GPU上设置模型
net = nn.DataParallel(net, device_ids=devices)
trainer = torch.optim.SGD(net.parameters(), lr)
loss = 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_iter:
trainer.zero_grad()
X, y = X.to(devices[0]), y.to(devices[0])
l = loss(net(X), y)
l.backward()
trainer.step()
timer.stop()
animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(net, test_iter),))
print(f'测试精度：{animator.Y[0][-1]:.2f}，{timer.avg():.1f}秒/轮，'
f'在{str(devices)}')
def train(net, num_gpus, 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)]
init_normal = nn.initializer.Normal(mean=0.0, std=0.01)
for i in net.sublayers():
if type(i) in [nn.Linear, nn.Conv2D]:
init_normal(i.weight)
# 在多个 GPU 上设置模型
net = paddle.DataParallel(net)
trainer = paddle.optimizer.SGD(parameters=net.parameters(), learning_rate=lr)
loss = 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_iter:
trainer.clear_grad()
X, y = paddle.to_tensor(X, place=devices[0]), paddle.to_tensor(y, place=devices[0])
l = loss(net(X), y)
l.backward()
trainer.step()
timer.stop()
animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(net, test_iter),))
print(f'测试精度：{animator.Y[0][-1]:.2f}, {timer.avg():.1f}秒/轮，'
f'在{str(devices)}')
接下来看看这在实践中是如何运作的。我们先在单个GPU上训练网络进行预热。
train(num_gpus=1, batch_size=256, lr=0.1)
测试精度：0.93，14.9秒/轮，在[gpu(0)]
train(net, num_gpus=1, batch_size=256, lr=0.1)
测试精度：0.93，12.2秒/轮，在[device(type='cuda', index=0)]
train(net, num_gpus=1, batch_size=256, lr=0.1)
测试精度：0.91, 13.1秒/轮，在[Place(gpu:0)]
接下来我们使用2个GPU进行训练。与 12.5节中评估的LeNet相比，ResNet-18的模型要复杂得多。这就是显示并行化优势的地方，计算所需时间明显大于同步参数需要的时间。因为并行化开销的相关性较小，因此这种操作提高了模型的可伸缩性。
train(num_gpus=2, batch_size=512, lr=0.2)
测试精度：0.92，7.3秒/轮，在[gpu(0), gpu(1)]
train(net, num_gpus=2, batch_size=512, lr=0.2)
测试精度：0.67，7.4秒/轮，在[device(type='cuda', index=0), device(type='cuda', index=1)]
12.6.4. 小结¶
Gluon通过提供一个上下文列表，为跨多个设备的模型初始化提供原语。
神经网络可以在（可找到数据的）单GPU上进行自动评估。
每台设备上的网络需要先初始化，然后再尝试访问该设备上的参数，否则会遇到错误。
优化算法在多个GPU上自动聚合。
神经网络可以在（可找到数据的）单GPU上进行自动评估。
每台设备上的网络需要先初始化，然后再尝试访问该设备上的参数，否则会遇到错误。
优化算法在多个GPU上自动聚合。
神经网络可以在（可找到数据的）单GPU上进行自动评估。
每台设备上的网络需要先初始化，然后再尝试访问该设备上的参数，否则会遇到错误。
优化算法在多个GPU上自动聚合。
12.6.5. 练习¶
本节使用ResNet-18，请尝试不同的迭代周期数、批量大小和学习率，以及使用更多的GPU进行计算。如果使用\(16\)个GPU（例如，在AWS p2.16xlarge实例上）尝试此操作，会发生什么？
有时候不同的设备提供了不同的计算能力，我们可以同时使用GPU和CPU，那应该如何分配工作？为什么？
如果去掉
npx.waitall()会怎样？该如何修改训练，以使并行操作最多有两个步骤重叠？
本节使用ResNet-18，请尝试不同的迭代周期数、批量大小和学习率，以及使用更多的GPU进行计算。如果使用\(16\)个GPU（例如，在AWS p2.16xlarge实例上）尝试此操作，会发生什么？
有时候不同的设备提供了不同的计算能力，我们可以同时使用GPU和CPU，那应该如何分配工作？为什么？
本节使用ResNet-18，请尝试不同的迭代周期数、批量大小和学习率，以及使用更多的GPU进行计算。如果使用\(16\)个GPU（例如，在AWS p2.16xlarge实例上）尝试此操作，会发生什么？
有时候不同的设备提供了不同的计算能力，我们可以同时使用GPU和CPU，那应该如何分配工作？为什么？