多 GPU 计算

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

In [1]:
!nvidia-smi
Mon Sep 17 12:32:46 2018
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 375.26                 Driver Version: 375.26                    |
|-------------------------------+----------------------+----------------------+
| 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           On   | 0000:00:1D.0     Off |                    0 |
| N/A   30C    P8    14W / 150W |      0MiB /  7612MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   1  Tesla M60           On   | 0000:00:1E.0     Off |                    0 |
| N/A   35C    P8    13W / 150W |      0MiB /  7612MiB |      0%      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的数据平行下的梯度计算。

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

In [2]:
import sys
sys.path.insert(0, '..')

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

定义模型

我们使用“卷积神经网络:LeNet”一节里介绍的 LeNet 来作为本节的样例模型。

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 之间同步数据的辅助函数。下面函数将模型参数复制到某个特定 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复制到mx.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 之间的数据。以下函数可以把各个 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)>]

给定一个批量的数据样本,以下函数可以划分它们并复制到各个 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 之间同步数据的辅助函数,例如split_and_loadallreduce

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
        net = lambda x: lenet(x, gpu_params[0])  # 在 GPU 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.5 sec, test acc: 0.10
epoch 2, time: 2.0 sec, test acc: 0.68
epoch 3, time: 2.0 sec, test acc: 0.76
epoch 4, time: 2.0 sec, test acc: 0.78

保持批量大小和学习率不变,将 GPU 改为 2,可以看到测试精度增加同前一致。特别是如果使用同样的初始模型参数,使用多个 GPU 应该得到跟单个 GPU 一样的结果。但由于每个 GPU 得到的批量大小减半了,它的计算效率(每秒能处理的样本数)变低了,而且加入了额外的通讯开销,所以我们并没有看到计算时间有显著下降。

In [13]:
train(num_gpus=2, batch_size=256, lr=0.2)
running on: [gpu(0), gpu(1)]
epoch 1, time: 2.1 sec, test acc: 0.10
epoch 2, time: 1.9 sec, test acc: 0.65
epoch 3, time: 1.9 sec, test acc: 0.75
epoch 4, time: 2.0 sec, test acc: 0.76

接下来我们将批量大小翻倍,这样每个 GPU 拿到同单 GPU 情况下一样的批量大小。同时,我们将学习率翻倍,希望能得到和单 GPU 下一样的测试精度递增速度。

In [14]:
train(num_gpus=2, batch_size=512, lr=0.4)
running on: [gpu(0), gpu(1)]
epoch 1, time: 1.4 sec, test acc: 0.10
epoch 2, time: 1.3 sec, test acc: 0.10
epoch 3, time: 1.3 sec, test acc: 0.34
epoch 4, time: 1.3 sec, test acc: 0.63

可以看到计算速度有了明显提升。这是因为我们将每个 GPU 的计算任务翻倍,这样得到了跟单 GPU 一样的计算效率。同时每个小批量的计算时间更长,但每次的通讯开销并没有增加,所以并行效率更高。而且在一个迭代周期里,通讯次数也减半。所有这些因素累加一起使得训练性能提升。图 8.2 示意了计算和通讯与总时间的关系。

小结

  • 我们可以使用数据并行更充分地利用多个 GPU 的计算资源,实现多 GPU 训练模型。
  • 给定超参数下的情况下,改变 GPU 个数不影响模型训练结果。
  • 总计算时间跟计算开销和通讯开销相关,这里批量大小是关键超参数。

练习

  • 在本节实验中,试一试不同的迭代周期、批量大小和学习率。
  • 将本节实验的模型预测部分改为用多 GPU 预测。
  • 测量每个批量纯计算时间和纯通讯时间。

扫码直达讨论区