多 GPU 计算¶

本节中我们将展示如何使用多个 GPU 计算，例如使用多个 GPU 训练同一个模型。正如你期望的那样，运行本节中的程序需要至少两块 GPU。事实上，一台机器上安装多块 GPU 很常见，这是因为主板上通常会有多个 PCIe 插槽。如果正确安装了 NVIDIA 驱动，我们可以通过nvidia-smi命令来查看当前计算机上的全部 GPU。

In [1]:

!nvidia-smi

Fri Dec 28 18:04:22 2018
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 396.37                 Driver Version: 396.37                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla M60           Off  | 00000000:00:1D.0 Off |                    0 |
| N/A   26C    P0    44W / 150W |      0MiB /  7618MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   1  Tesla M60           Off  | 00000000:00:1E.0 Off |                    0 |
| N/A   34C    P0    41W / 150W |      0MiB /  7618MiB |     96%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

“自动并行计算”一节介绍过，大部分的运算可以使用所有的 CPU 的全部计算资源，或者单个 GPU 的全部计算资源。但如果使用多个 GPU 训练模型，我们仍然需要实现相应的算法。这些算法中最常用的叫做数据并行。

数据并行¶

数据并行目前是深度学习里使用最广泛的将模型训练任务划分到多个 GPU 的办法。回忆一下我们在“小批量随机梯度下降”一节中介绍的使用优化算法训练模型的过程。下面我们就以小批量随机梯度下降为例来介绍数据并行是如何工作的。

假设一台机器上有 \(k\) 个 GPU。给定需要训练的模型，每个 GPU 将分别独立维护一份完整的模型参数。在模型训练的任意一次迭代中，给定一个随机小批量，我们将该批量中的样本划分成 \(k\) 份并分给每个 GPU 一份。然后，每个 GPU 将根据自己所分到的小批量子集和自己所维护的模型参数分别计算模型参数的本地梯度。接下来，我们把 \(k\) 个 GPU 上的本地梯度相加，便得到当前的小批量随机梯度。之后，每个 GPU 都使用这个小批量随机梯度分别更新自己所维护的那一份完整的模型参数。图 8.1 描绘了使用两个 GPU 的数据并行下的小批量随机梯度的计算。

使用两个GPU的数据并行下的小批量随机梯度的计算。

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

In [2]:

import gluonbook as gb
import mxnet as mx
from mxnet import autograd, nd
from mxnet.gluon import loss as gloss
import time

定义模型¶

我们使用“卷积神经网络（LeNet）”一节里介绍的 LeNet 来作为本节的样例模型。这里的模型实现部分只用到了 NDArray。

In [3]:

# 初始化模型参数。
scale = 0.01
W1 = nd.random.normal(scale=scale, shape=(20, 1, 3, 3))
b1 = nd.zeros(shape=20)
W2 = nd.random.normal(scale=scale, shape=(50, 20, 5, 5))
b2 = nd.zeros(shape=50)
W3 = nd.random.normal(scale=scale, shape=(800, 128))
b3 = nd.zeros(shape=128)
W4 = nd.random.normal(scale=scale, shape=(128, 10))
b4 = nd.zeros(shape=10)
params = [W1, b1, W2, b2, W3, b3, W4, b4]

# 定义模型。
def lenet(X, params):
    h1_conv = nd.Convolution(data=X, weight=params[0], bias=params[1],
                             kernel=(3, 3), num_filter=20)
    h1_activation = nd.relu(h1_conv)
    h1 = nd.Pooling(data=h1_activation, pool_type='avg', kernel=(2, 2),
                    stride=(2, 2))
    h2_conv = nd.Convolution(data=h1, weight=params[2], bias=params[3],
                             kernel=(5, 5), num_filter=50)
    h2_activation = nd.relu(h2_conv)
    h2 = nd.Pooling(data=h2_activation, pool_type='avg', kernel=(2, 2),
                    stride=(2, 2))
    h2 = nd.flatten(h2)
    h3_linear = nd.dot(h2, params[4]) + params[5]
    h3 = nd.relu(h3_linear)
    y_hat = nd.dot(h3, params[6]) + params[7]
    return y_hat

# 交叉熵损失函数。
loss = gloss.SoftmaxCrossEntropyLoss()

多 GPU 之间同步数据¶

我们需要实现一些多 GPU 之间同步数据的辅助函数。下面的get_params函数将模型参数复制到某个特定 GPU 并初始化梯度。

In [4]:

def get_params(params, ctx):
    new_params = [p.copyto(ctx) for p in params]
    for p in new_params:
        p.attach_grad()
    return new_params

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

In [5]:

new_params = get_params(params, mx.gpu(0))
print('b1 weight:', new_params[1])
print('b1 grad:', new_params[1].grad)

b1 weight:
[0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
<NDArray 20 @gpu(0)>
b1 grad:
[0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
<NDArray 20 @gpu(0)>

给定分布在多个 GPU 之间的数据。以下的allreduce函数可以把各个 GPU 上的数据加起来，然后再广播到所有的 GPU 上。

In [6]:

def allreduce(data):
    for i in range(1, len(data)):
        data[0][:] += data[i].copyto(data[0].context)
    for i in range(1, len(data)):
        data[0].copyto(data[i])

简单测试一下allreduce函数。

In [7]:

data = [nd.ones((1, 2), ctx=mx.gpu(i)) * (i + 1) for i in range(2)]
print('before allreduce:', data)
allreduce(data)
print('after allreduce:', data)

before allreduce: [
[[1. 1.]]
<NDArray 1x2 @gpu(0)>,
[[2. 2.]]
<NDArray 1x2 @gpu(1)>]
after allreduce: [
[[3. 3.]]
<NDArray 1x2 @gpu(0)>,
[[3. 3.]]
<NDArray 1x2 @gpu(1)>]

给定一个批量的数据样本，以下的split_and_load函数可以划分它们并复制到各个 GPU 上。

In [8]:

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

让我们试着用split_and_load函数将 6 个数据样本平均分给 2 个 GPU。

In [9]:

batch = nd.arange(24).reshape((6, 4))
ctx = [mx.gpu(0), mx.gpu(1)]
splitted = split_and_load(batch, ctx)
print('input: ', batch)
print('load into', ctx)
print('output:', splitted)

input:
[[ 0.  1.  2.  3.]
 [ 4.  5.  6.  7.]
 [ 8.  9. 10. 11.]
 [12. 13. 14. 15.]
 [16. 17. 18. 19.]
 [20. 21. 22. 23.]]
<NDArray 6x4 @cpu(0)>
load into [gpu(0), gpu(1)]
output: [
[[ 0.  1.  2.  3.]
 [ 4.  5.  6.  7.]
 [ 8.  9. 10. 11.]]
<NDArray 3x4 @gpu(0)>,
[[12. 13. 14. 15.]
 [16. 17. 18. 19.]
 [20. 21. 22. 23.]]
<NDArray 3x4 @gpu(1)>]

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

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

In [10]:

def train_batch(X, y, gpu_params, ctx, lr):
    # 当 ctx 包含多个 GPU 时，划分小批量数据样本并复制到各个 GPU 上。
    gpu_Xs, gpu_ys = split_and_load(X, ctx), split_and_load(y, ctx)
    with autograd.record():  # 在各个 GPU 上分别计算损失。
        ls = [loss(lenet(gpu_X, gpu_W), gpu_y)
              for gpu_X, gpu_y, gpu_W in zip(gpu_Xs, gpu_ys, gpu_params)]
    for l in ls:  # 在各个 GPU 上分别反向传播。
        l.backward()
    # 把各个 GPU 上的梯度加起来，然后再广播到所有 GPU 上。
    for i in range(len(gpu_params[0])):
        allreduce([gpu_params[c][i].grad for c in range(len(ctx))])
    for param in gpu_params:  # 在各个 GPU 上分别更新模型参数。
        gb.sgd(param, lr, X.shape[0])  # 这里使用了完整批量大小。

训练函数¶

现在我们可以定义训练函数。这里的训练函数和之前章节里的训练函数稍有不同。例如，在这里我们需要依据数据并行将完整的模型参数复制到多个 GPU 上，并在每次迭代时对单个小批量上进行多 GPU 训练。

In [11]:

def train(num_gpus, batch_size, lr):
    train_iter, test_iter = gb.load_data_fashion_mnist(batch_size)
    ctx = [mx.gpu(i) for i in range(num_gpus)]
    print('running on:', ctx)
    # 将模型参数复制到 num_gpus 个 GPU 上。
    gpu_params = [get_params(params, c) for c in ctx]
    for epoch in range(4):
        start = time.time()
        for X, y in train_iter:
            # 对单个小批量进行多 GPU 训练。
            train_batch(X, y, gpu_params, ctx, lr)
            nd.waitall()
        train_time = time.time() - start

        def net(x):  # 在 GPU 0 上验证模型。
            return lenet(x, gpu_params[0])

        test_acc = gb.evaluate_accuracy(test_iter, net, ctx[0])
        print('epoch %d, time: %.1f sec, test acc: %.2f'
              % (epoch + 1, train_time, test_acc))

多 GPU 训练实验¶

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

In [12]:

train(num_gpus=1, batch_size=256, lr=0.2)

running on: [gpu(0)]
epoch 1, time: 2.7 sec, test acc: 0.10
epoch 2, time: 2.2 sec, test acc: 0.69
epoch 3, time: 2.2 sec, test acc: 0.78
epoch 4, time: 2.2 sec, test acc: 0.77

保持批量大小和学习率不变，将使用的 GPU 数改为 2，可以看到测试精度的提升同上一个实验中的结果大体相当。由于额外的通讯开销，我们并没有看到训练时间的显著降低。

In [13]:

train(num_gpus=2, batch_size=256, lr=0.2)

running on: [gpu(0), gpu(1)]
epoch 1, time: 2.6 sec, test acc: 0.17
epoch 2, time: 2.4 sec, test acc: 0.65
epoch 3, time: 2.2 sec, test acc: 0.68
epoch 4, time: 2.2 sec, test acc: 0.76

小结¶

我们可以使用数据并行更充分地利用多个 GPU 的计算资源，实现多 GPU 训练模型。
给定超参数的情况下，改变 GPU 个数时模型的训练精度大体相当。

练习¶

在多 GPU 训练实验中，使用 2 个 GPU 训练并将batch_size翻倍至 512，训练时间有何变化？如果希望测试精度与单 GPU 训练中的结果相当，学习率应如何调节？
将实验的模型预测部分改为用多 GPU 预测。