多 GPU 计算

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

In [1]:
!nvidia-smi
Wed Jul 18 06:57:05 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   34C    P0    37W / 150W |    298MiB /  7612MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   1  Tesla M60           On   | 0000:00:1E.0     Off |                    0 |
| N/A   43C    P0    37W / 150W |    289MiB /  7612MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID  Type  Process name                               Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

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

数据并行

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

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

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

In [2]:
import sys
sys.path.append('..')
import gluonbook as gb
import mxnet as mx
from mxnet import autograd, nd
from mxnet.gluon import loss as gloss
from time 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 = split_and_load(X, ctx)
    gpu_ys = split_and_load(y, ctx)
    # 在各个 GPU 上计算损失。
    with autograd.record():
        ls = [loss(lenet(gpu_X, gpu_W), gpu_y)
              for gpu_X, gpu_y, gpu_W in zip(gpu_Xs, gpu_ys, gpu_params)]
    # 在各个 GPU 上反向传播。
    for l in ls:
        l.backward()
    # 把各个 GPU 上的梯度加起来,然后再广播到所有 GPU 上。
    for i in range(len(gpu_params[0])):
        allreduce([gpu_params[c][i].grad for c in range(len(ctx))])
    # 在各个 GPU 上更新自己维护的那一份完整的模型参数。
    for param in gpu_params:
        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(1, 6):
        start = time()
        for X, y in train_iter:
            # 对单个小批量上进行多 GPU 训练。
            train_batch(X, y, gpu_params, ctx, lr)
        nd.waitall()
        print('epoch %d, time: %.1f sec' % (epoch, time() - start))
        # 在 GPU0 上验证模型。
        net = lambda x: lenet(x, gpu_params[0])
        test_acc = gb.evaluate_accuracy(test_iter, net, ctx[0])
        print('validation accuracy: %.4f' % test_acc)

我们先使用一个 GPU 来训练。

In [12]:
train(num_gpus=1, batch_size=256, lr=0.3)
running on: [gpu(0)]
epoch 1, time: 2.2 sec
validation accuracy: 0.1611
epoch 2, time: 1.8 sec
validation accuracy: 0.7298
epoch 3, time: 1.8 sec
validation accuracy: 0.7414
epoch 4, time: 1.8 sec
validation accuracy: 0.7795
epoch 5, time: 1.8 sec
validation accuracy: 0.8235

接下来,我们先使用 2 个 GPU 来训练。我们将批量大小也增加一倍,以使得 GPU 的计算资源能够得到较充分利用。

In [13]:
train(num_gpus=2, batch_size=512, lr=0.3)
running on: [gpu(0), gpu(1)]
epoch 1, time: 1.4 sec
validation accuracy: 0.1000
epoch 2, time: 1.2 sec
validation accuracy: 0.2138
epoch 3, time: 1.2 sec
validation accuracy: 0.6997
epoch 4, time: 1.3 sec
validation accuracy: 0.7116
epoch 5, time: 1.3 sec
validation accuracy: 0.7712

由于批量大小增加了一倍,每个迭代周期的迭代次数减小了一半。因此,我们观察到每个迭代周期的耗时比单 GPU 训练时少了近一半。但由于总体迭代次数的减少,模型在验证数据集上的精度略有下降。这很可能是由于训练不够充分造成的。因此,多 GPU 训练时,我们可以适当增加迭代周期使训练较充分。

小结

  • 我们可以使用数据并行更充分地利用多个 GPU 的计算资源,实现多 GPU 训练模型。

练习

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

扫码直达 讨论区