深度学习：如何使用多块 GPU 计算？

本节中咱们将展现如何使用多块 GPU 计算，例如，使用多块GPU 训练同一个模型。正如所指望的那样，运行本节中的程序须要至少2块 GPU。事实上，一台机器上安装多块 GPU 很常见，这是由于主板上一般会有多个 PCIe 插槽。若是正确安装了 NVIDIA 驱动，咱们能够经过nvidia-smi命令来查看当前计算机上的所有 GPU。算法

1. In [1]: !nvidia-smi
2.  
3. Mon Feb 25 19:19:54 2019
4. +-----------------------------------------------------------------------------+
5. | NVIDIA-SMI 384.111 Driver Version: 384.111 |
6. |-------------------------------+----------------------+----------------------+
7. | GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
8. | Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
9. |===============================+======================+======================|
10. | 0 Tesla V100-SXM2… On | 00000000:00:1B.0 Off | 0 |
11. | N/A 46C P0 38W / 300W | 0MiB / 16152MiB | 0% Default |
12. +-------------------------------+----------------------+----------------------+
13. | 1 Tesla V100-SXM2… On | 00000000:00:1C.0 Off | 0 |
14. | N/A 44C P0 39W / 300W | 0MiB / 16152MiB | 0% Default |
15. +-------------------------------+----------------------+----------------------+
16. | 2 Tesla V100-SXM2… On | 00000000:00:1D.0 Off | 0 |
17. | N/A 42C P0 39W / 300W | 0MiB / 16152MiB | 0% Default |
18. +-------------------------------+----------------------+----------------------+
19. | 3 Tesla V100-SXM2… On | 00000000:00:1E.0 Off | 0 |
20. | N/A 45C P0 43W / 300W | 0MiB / 16152MiB | 0% Default |
21. +-------------------------------+----------------------+----------------------+
22.  
23. +-----------------------------------------------------------------------------+
24. | Processes: GPU Memory |
25. | GPU PID Type Process name Usage |
26. |=============================================================================|
27. | No running processes found |
28. +-----------------------------------------------------------------------------+

8.3节介绍过，大部分运算可使用全部的 CPU 的所有计算资源，或者单块GPU 的所有计算资源。但若是使用多块GPU 训练模型，咱们仍然须要实现相应的算法。这些算法中最经常使用的叫做数据并行。编程

8.4.1　数据并行

数据并行目前是深度学习里使用最普遍的将模型训练任务划分到多块GPU 的方法。回忆一下咱们在7.3节中介绍的使用优化算法训练模型的过程。下面咱们就以小批量随机梯度降低为例来介绍数据并行是如何工做的。dom

假设一台机器上有k块GPU。给定须要训练的模型，每块GPU及其相应的显存将分别独立维护一份完整的模型参数。在模型训练的任意一次迭代中，给定一个随机小批量，咱们将该批量中的样本划分红k份并分给每块显卡的显存一份。而后，每块GPU 将根据相应显存所分到的小批量子集和所维护的模型参数分别计算模型参数的本地梯度。接下来，咱们把k块显卡的显存上的本地梯度相加，便获得当前的小批量随机梯度。以后，每块GPU 都使用这个小批量随机梯度分别更新相应显存所维护的那一份完整的模型参数。图8-1 描绘了使用2块GPU 的数据并行下的小批量随机梯度的计算。机器学习

图8-1　使用2块GPU的数据并行下的小批量随机梯度的计算ide

为了从零开始实现多 GPU 训练中的数据并行，让咱们先导入须要的包或模块。函数

1. In [2]: import d2lzh as d2l
2. import mxnet as mx
3. from mxnet import autograd, nd
4. from mxnet.gluon import loss as gloss
5. import time

8.4.2　定义模型

咱们使用5.5节里介绍的LeNet来做为本节的样例模型。这里的模型实现部分只用到了 NDArray。性能

1. In [3]: # 初始化模型参数
2. scale = 0.01
3. W1 = nd.random.normal(scale=scale, shape=(20, 1, 3, 3))
4. b1 = nd.zeros(shape=20)
5. W2 = nd.random.normal(scale=scale, shape=(50, 20, 5, 5))
6. b2 = nd.zeros(shape=50)
7. W3 = nd.random.normal(scale=scale, shape=(800, 128))
8. b3 = nd.zeros(shape=128)
9. W4 = nd.random.normal(scale=scale, shape=(128, 10))
10. b4 = nd.zeros(shape=10)
11. params = [W1, b1, W2, b2, W3, b3, W4, b4]
12.  
13. # 定义模型
14. def lenet(X, params):
15. h1_conv = nd.Convolution(data=X, weight=params[0], bias=params[1],
16. kernel=(3, 3), num_f ilter=20)
17. h1_activation = nd.relu(h1_conv)
18. h1 = nd.Pooling(data=h1_activation, pool_type='avg', kernel=(2, 2),
19. stride=(2, 2))
20. h2_conv = nd.Convolution(data=h1, weight=params[2], bias=params[3],
21. kernel=(5, 5), num_f ilter=50)
22. h2_activation = nd.relu(h2_conv)
23. h2 = nd.Pooling(data=h2_activation, pool_type='avg', kernel=(2, 2),
24. stride=(2, 2))
25. h2 = nd.f latten(h2)
26. h3_linear = nd.dot(h2, params[4]) + params[5]
27. h3 = nd.relu(h3_linear)
28. y_hat = nd.dot(h3, params[6]) + params[7]
29. return y_hat
30.  
31. # 交叉熵损失函数
32. loss = gloss.SoftmaxCrossEntropyLoss()

8.4.3　多GPU之间同步数据

咱们须要实现一些多GPU之间同步数据的辅助函数。下面的get_params函数将模型参数复制到某块显卡的显存并初始化梯度。学习

1. In [4]: def get_params(params, ctx):
2. new_params = [p.copyto(ctx) for p in params]
3. for p in new_params:
4. p.attach_grad()
5. return new_params

尝试把模型参数params复制到gpu(0)上。测试

1. In [5]: new_params = get_params(params, mx.gpu(0))
2. print('b1 weight:', new_params[1])
3. print('b1 grad:', new_params[1].grad)
4.  
5. b1 weight:
6. [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
7. <NDArray 20 @gpu(0)>
8. b1 grad:
9. [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
10. <NDArray 20 @gpu(0)>

给定分布在多块显卡的显存之间的数据。下面的allreduce函数能够把各块显卡的显存上的数据加起来，而后再广播到全部的显存上。优化

1. In [6]: def allreduce(data):
2. for i in range(1, len(data)):
3. data[0][:] += data[i].copyto(data[0].context)
4. for i in range(1, len(data)):
5. data[0].copyto(data[i])

简单测试一下allreduce函数。

1. In [7]: data = [nd.ones((1, 2), ctx=mx.gpu(i)) * (i + 1) for i in range(2)]
2. print('before allreduce:', data)
3. allreduce(data)
4. print('after allreduce:', data)
5.  
6. before allreduce: [
7. [[1. 1.]]
8. <NDArray 1x2 @gpu(0)>,
9. [[2. 2.]]
10. <NDArray 1x2 @gpu(1)>]
11. after allreduce: [
12. [[3. 3.]]
13. <NDArray 1x2 @gpu(0)>,
14. [[3. 3.]]
15. <NDArray 1x2 @gpu(1)>]

给定一个批量的数据样本，下面的split_and_load函数能够将其划分并复制到各块显卡的显存上。

1. In [8]: def split_and_load(data, ctx):
2. n, k = data.shape[0], len(ctx)
3. m = n // k # 简单起见, 假设能够整除
4. assert m * k == n, '# examples is not divided by # devices.'
5. return [data[i * m: (i + 1) * m].as_in_context(ctx[i]) for i in range(k)]

让咱们试着用split_and_load函数将6个数据样本平均分给2块显卡的显存。

1. In [9]: batch = nd.arange(24).reshape((6, 4))
2. ctx = [mx.gpu(0), mx.gpu(1)]
3. splitted = split_and_load(batch, ctx)
4. print('input: ', batch)
5. print('load into', ctx)
6. print('output:', splitted)
7.  
8. input:
9. [[ 0. 1. 2. 3.]
10. [ 4. 5. 6. 7.]
11. [ 8. 9. 10. 11.]
12. [12. 13. 14. 15.]
13. [16. 17. 18. 19.]
14. [20. 21. 22. 23.]]
15. <NDArray 6x4 @cpu(0)>
16. load into [gpu(0), gpu(1)]
17. output: [
18. [[ 0. 1. 2. 3.]
19. [ 4. 5. 6. 7.]
20. [ 8. 9. 10. 11.]]
21. <NDArray 3x4 @gpu(0)>,
22. [[12. 13. 14. 15.]
23. [16. 17. 18. 19.]
24. [20. 21. 22. 23.]]
25. <NDArray 3x4 @gpu(1)>]

8.4.4　单个小批量上的多GPU训练

如今咱们能够实现单个小批量上的多GPU训练了。它的实现主要依据本节介绍的数据并行方法。咱们将使用刚刚定义的多GPU之间同步数据的辅助函数allreduce和split_and_load。

1. In [10]: def train_batch(X, y, gpu_params, ctx, lr):
2. # 当ctx包含多块GPU及相应的显存时, 将小批量数据样本划分并复制到各个显存上
3. gpu_Xs, gpu_ys = split_and_load(X, ctx), split_and_load(y, ctx)
4. with autograd.record(): # 在各块GPU上分别计算损失
5. ls = [loss(lenet(gpu_X, gpu_W), gpu_y)
6. for gpu_X, gpu_y, gpu_W in zip(gpu_Xs, gpu_ys, gpu_params)]
7. for l in ls: # 在各块GPU上分别反向传播
8. l.backward()
9. # 把各块显卡的显存上的梯度加起来, 而后广播到全部显存上
10. for i in range(len(gpu_params[0])):
11. allreduce([gpu_params[c][i].grad for c in range(len(ctx))])
12. for param in gpu_params: # 在各块显卡的显存上分别更新模型参数
13. d2l.sgd(param, lr, X.shape[0]) # 这里使用了完整批量大小

8.4.5　定义训练函数

如今咱们能够定义训练函数了。这里的训练函数和3.6节定义的训练函数train_ch3有所不一样。值得强调的是，在这里咱们须要依据数据并行将完整的模型参数复制到多块显卡的显存上，并在每次迭代时对单个小批量进行多GPU训练。

1. In [11]: def train(num_gpus, batch_size, lr):
2. train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
3. ctx = [mx.gpu(i) for i in range(num_gpus)]
4. print('running on:', ctx)
5. # 将模型参数复制到num_gpus块显卡的显存上
6. gpu_params = [get_params(params, c) for c in ctx]
7. for epoch in range(4):
8. start = time.time()
9. for X, y in train_iter:
10. # 对单个小批量进行多GPU训练
11. train_batch(X, y, gpu_params, ctx, lr)
12. nd.waitall()
13. train_time = time.time() - start
14.  
15. def net(x): # 在gpu(0)上验证模型
16. return lenet(x, gpu_params[0])
17. test_acc = d2l.evaluate_accuracy(test_iter, net, ctx[0])
18. print('epoch %d, time %.1f sec, test acc %.2f'
19. % (epoch + 1, train_time, test_acc))

8.4.6　多GPU训练实验

让咱们先从单 GPU训练开始。设批量大小为 256，学习率为 0.2。

1. In [12]: train(num_gpus=1, batch_size=256, lr=0.2)
2.  
3. running on: [gpu(0)]
4. epoch 1, time 1.7 sec, test acc 0.10
5. epoch 2, time 1.6 sec, test acc 0.69
6. epoch 3, time 1.5 sec, test acc 0.75
7. epoch 4, time 1.6 sec, test acc 0.79

保持批量大小和学习率不变，将使用的GPU数量改成2。能够看到，测试准确率的提高同上一个实验中的结果大致至关。由于有额外的通讯开销，因此咱们并无看到训练时间的显著下降。所以，咱们将在8.5节实验计算更加复杂的模型。

1. In [13]: train(num_gpus=2, batch_size=256, lr=0.2)
2.  
3. running on: [gpu(0), gpu(1)]
4. epoch 1, time 2.5 sec, test acc 0.10
5. epoch 2, time 2.3 sec, test acc 0.64
6. epoch 3, time 2.4 sec, test acc 0.68
7. epoch 4, time 2.6 sec, test acc 0.78

小结

• 可使用数据并行更充分地利用多块GPU 的计算资源，实现多 GPU 训练模型。

• 给定超参数的状况下，改变 GPU数量时模型的准确率大致至关。

练习

（1）在多 GPU 训练实验中，使用2块GPU 训练并将batch_size翻倍至512，训练时间有何变化？若是但愿测试准确率与单GPU训练中的结果至关，学习率应如何调节？

（2）将实验的模型预测部分改成用多 GPU 预测。

8.5　多GPU计算的简洁实现

在 Gluon 中，咱们能够很方便地使用数据并行进行多GPU计算。例如，咱们并不须要本身实现8.4节里介绍的多 GPU 之间同步数据的辅助函数。

首先导入本节实验所需的包或模块。运行本节中的程序须要至少2块GPU。

1. In [1]: import d2lzh as d2l
2. import mxnet as mx
3. from mxnet import autograd, gluon, init, nd
4. from mxnet.gluon import loss as gloss, nn, utils as gutils
5. import time

8.5.1　多GPU上初始化模型参数

咱们使用 ResNet-18做为本节的样例模型。因为本节的输入图像使用原尺寸（未放大），这里的模型构造与5.11节中的 ResNet-18 构造稍有不一样。这里的模型在一开始使用了较小的卷积核、步幅和填充，并去掉了最大池化层。

1. In [2]: def resnet18(num_classes): # 本函数已保存在d2lzh包中方便之后使用
2. def resnet_block(num_channels, num_residuals, f irst_block=False):
3. blk = nn.Sequential()
4. for i in range(num_residuals):
5. if i == 0 and not f irst_block:
6. blk.add(d2l.Residual(
7. num_channels, use_1x1conv=True, strides=2))
8. else:
9. blk.add(d2l.Residual(num_channels))
10. return blk
11.  
12. net = nn.Sequential()
13. # 这里使用了较小的卷积核、步幅和填充, 并去掉了最大池化层
14. net.add(nn.Conv2D(64, kernel_size=3, strides=1, padding=1),
15. nn.BatchNorm(), nn.Activation('relu'))
16. net.add(resnet_block(64, 2, f irst_block=True),
17. resnet_block(128, 2),
18. resnet_block(256, 2),
19. resnet_block(512, 2))
20. net.add(nn.GlobalAvgPool2D(), nn.Dense(num_classes))
21. return net
22.  
23. net = resnet18(10)

以前咱们介绍了如何使用initialize函数的ctx参数在内存或单块显卡的显存上初始化模型参数。事实上，ctx能够接受一系列的 CPU及内存和GPU及相应的显存，从而使初始化好的模型参数复制到ctx里全部的内存和显存上。

1. In [3]: ctx = [mx.gpu(0), mx.gpu(1)]
2. net.initialize(init=init.Normal(sigma=0.01), ctx=ctx)

Gluon提供了上一节中实现的split_and_load函数。它能够划分一个小批量的数据样本并复制到各个内存或显存上。以后，根据输入数据所在的内存或显存，模型计算会相应地使用CPU或相同显卡上的GPU。

1. In [4]: x = nd.random.uniform(shape=(4, 1, 28, 28))
2. gpu_x = gutils.split_and_load(x, ctx)
3. net(gpu_x[0]), net(gpu_x[1])
4.  
5. Out[4]: (
6. [[ 5.4814936e-06 -8.3371094e-07 -1.6316770e-06 -6.3674099e-07
7. -3.8216162e-06 -2.3514044e-06 -2.5469599e-06 -9.4784696e-08
8. -6.9033558e-07 2.5756231e-06]
9. [ 5.4710872e-06 -9.4246496e-07 -1.0494070e-06 9.8081841e-08
10. -3.3251815e-06 -2.4862918e-06 -3.3642798e-06 1.0455864e-07
11. -6.1001344e-07 2.0327841e-06]]
12. <NDArray 2x10 @gpu(0)>,
13. [[ 5.6176345e-06 -1.2837586e-06 -1.4605541e-06 1.8302967e-07
14. -3.5511653e-06 -2.4371013e-06 -3.5731798e-06 -3.0974860e-07
15. -1.1016571e-06 1.8909889e-06]
16. [ 5.1418697e-06 -1.3729932e-06 -1.1520088e-06 1.1507450e-07
17. -3.7372811e-06 -2.8289724e-06 -3.6477197e-06 1.5781629e-07
18. -6.0733043e-07 1.9712013e-06]]
19. <NDArray 2x10 @gpu(1)>)

如今，咱们能够访问已初始化好的模型参数值了。须要注意的是，默认状况下weight.data()会返回内存上的参数值。由于咱们指定了2块GPU来初始化模型参数，因此须要指定显存来访问参数值。咱们看到，相同参数在不一样显卡的显存上的值同样。

1. In [5]: weight = net[0].params.get('weight')
2.  
3. try:
4. weight.data()
5. except RuntimeError:
6. print('not initialized on', mx.cpu())
7. weight.data(ctx[0])[0], weight.data(ctx[1])[0]
8.  
9. not initialized on cpu(0)
10.  
11. Out[5]: (
12. [[[-0.01473444 -0.01073093 -0.01042483]
13. [-0.01327885 -0.01474966 -0.00524142]
14. [ 0.01266256 0.00895064 -0.00601594]]]
15. <NDArray 1x3x3 @gpu(0)>,
16. [[[-0.01473444 -0.01073093 -0.01042483]
17. [-0.01327885 -0.01474966 -0.00524142]
18. [ 0.01266256 0.00895064 -0.00601594]]]
19. <NDArray 1x3x3 @gpu(1)>)

8.5.2　多GPU训练模型

当使用多块GPU来训练模型时，Trainer实例会自动作数据并行，例如，划分小批量数据样本并复制到各块显卡的显存上，以及对各块显卡的显存上的梯度求和再广播到全部显存上。这样，咱们就能够很方便地实现训练函数了。

1. In [6]: def train(num_gpus, batch_size, lr):
2. train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
3. ctx = [mx.gpu(i) for i in range(num_gpus)]
4. print('running on:', ctx)
5. net.initialize(init=init.Normal(sigma=0.01), ctx=ctx, force_reinit=True)
6. trainer = gluon.Trainer(
7. net.collect_params(), 'sgd', {'learning_rate': lr})
8. loss = gloss.SoftmaxCrossEntropyLoss()
9. for epoch in range(4):
10. start = time.time()
11. for X, y in train_iter:
12. gpu_Xs = gutils.split_and_load(X, ctx)
13. gpu_ys = gutils.split_and_load(y, ctx)
14. with autograd.record():
15. ls = [loss(net(gpu_X), gpu_y)
16. for gpu_X, gpu_y in zip(gpu_Xs, gpu_ys)]
17. for l in ls:
18. l.backward()
19. trainer.step(batch_size)
20. nd.waitall()
21. train_time = time.time() - start
22. test_acc = d2l.evaluate_accuracy(test_iter, net, ctx[0])
23. print('epoch %d, time %.1f sec, test acc %.2f' % (
24. epoch + 1, train_time, test_acc))

首先在单块GPU上训练模型。

1. In [7]: train(num_gpus=1, batch_size=256, lr=0.1)
2.  
3. running on: [gpu(0)]
4. epoch 1, time 14.9 sec, test acc 0.87
5. epoch 2, time 13.5 sec, test acc 0.90
6. epoch 3, time 13.6 sec, test acc 0.92
7. epoch 4, time 13.6 sec, test acc 0.91

而后尝试在2块GPU上训练模型。与8.4节使用的LeNet 相比，ResNet-18 的计算更加复杂，通讯时间比计算时间更短，所以 ResNet-18 的并行计算所得到的性能提高更佳。

1. In [8]: train(num_gpus=2, batch_size=512, lr=0.2)
2.  
3. running on: [gpu(0), gpu(1)]
4. epoch 1, time 7.8 sec, test acc 0.81
5. epoch 2, time 7.0 sec, test acc 0.87
6. epoch 3, time 7.0 sec, test acc 0.89
7. epoch 4, time 7.0 sec, test acc 0.91

小结

• 在Gluon中，能够很方便地进行多GPU计算，例如，在多GPU及相应的显存上初始化模型参数和训练模型。

本文截选自《动手学深度学习》，阿斯顿·张（Aston Zhang），李沐（Mu Li），[美] 扎卡里·C. 立顿 等 著。

本书面向但愿了解深度学习，特别是对实际使用深度学习感兴趣的大学生、工程师和研究人员。本书不要求读者有任何深度学习或者机器学习的背景知识，读者只需具有基本的数学和编程知识，如基础的线性代数、微分、几率及Python编程知识。本书的附录中提供了书中涉及的主要数学知识，供读者参考。

本书的英文版Dive into Deep Learning是加州大学伯克利分校2019年春学期“Introduction to Deep Learning”（深度学习导论）课程的教材。截至2019年春学期，本书中的内容已被全球15 所知名大学用于教学。本书的学习社区、免费教学资源（课件、教学视频、更多习题等），以及用于本书学习和教学的免费计算资源（仅限学生和老师）的申请方法在本书配套网站zh.d2l.ai上发布。读者在阅读本书的过程当中，若是对书中某节内容有疑惑，也能够扫一扫书中对应的二维码寻求帮助。