多GPU计算——使用Gluon

在Gluon中,我们可以很方便地使用数据并行进行多GPU计算。比方说,我们并不需要自己实现“多GPU计算——从零开始”一节里介绍的多GPU之间同步数据的辅助函数。

先导入本节实验需要的包或模块。同上一节,运行本节中的程序需要至少两块GPU。

In [1]:
import sys
sys.path.append('..')
import gluonbook as gb
import mxnet as mx
from mxnet import autograd, gluon, init, nd
from mxnet.gluon import loss as gloss, utils as gutils
from time import time

多GPU上初始化模型参数

我们使用ResNet-18来作为本节的样例模型。

In [2]:
net = gb.resnet18(10)

之前我们介绍了如何使用initialize函数的ctx参数在CPU或单个GPU上初始化模型参数。事实上,ctx可以接受一系列的CPU/GPU,从而使初始化好的模型参数复制到ctx里所有的CPU/GPU上。

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

Gluon提供了上一节中实现的split_and_load函数。它可以划分一个小批量的数据样本并复制到各个CPU/GPU上。之后,根据输入数据所在的CPU/GPU,模型计算会发生在相同的CPU/GPU上。

In [4]:
x = nd.random.uniform(shape=(4, 1, 28, 28))
gpu_x = gutils.split_and_load(x, ctx)
net(gpu_x[0]), net(gpu_x[1])
Out[4]:
(
 [[  3.22364103e-05  -1.77604452e-05  -4.53473040e-05  -4.43779390e-06
     4.11345136e-05  -1.12404132e-05  -3.77615615e-05   4.36321279e-05
     2.28628483e-06  -1.48372719e-05]
  [  3.10487958e-05  -1.89744533e-05  -4.50448642e-05  -5.33929597e-06
     4.11167130e-05  -1.23149039e-05  -3.99470700e-05   4.46575323e-05
     2.58294403e-06  -1.34151524e-05]]
 <NDArray 2x10 @gpu(0)>,
 [[  2.98842060e-05  -1.86083871e-05  -4.19585049e-05  -5.71156716e-06
     4.02533478e-05  -1.19904144e-05  -3.85621279e-05   4.39719333e-05
     6.29562464e-08  -1.58273233e-05]
  [  3.08366398e-05  -1.81654450e-05  -4.28299536e-05  -5.13964960e-06
     3.97496297e-05  -1.23091431e-05  -3.68527180e-05   4.15482536e-05
     3.11583290e-06  -1.18256366e-05]]
 <NDArray 2x10 @gpu(1)>)

回忆一下“模型参数的延后初始化”一节中介绍的延后的初始化。现在,我们可以通过data访问初始化好的模型参数值了。需要注意的是,默认下weight.data()会返回CPU上的参数值。由于我们指定了2个GPU来初始化模型参数,我们需要指定GPU访问。我们看到,相同参数在不同的GPU上的值一样。

In [5]:
weight = net[1].params.get('weight')
try:
    weight.data()
except:
    print('not initialized on', mx.cpu())
weight.data(ctx[0])[0], weight.data(ctx[1])[0]
not initialized on cpu(0)
Out[5]:
(
 [[[-0.01473444 -0.01073093 -0.01042483]
   [-0.01327885 -0.01474966 -0.00524142]
   [ 0.01266256  0.00895064 -0.00601594]]]
 <NDArray 1x3x3 @gpu(0)>,
 [[[-0.01473444 -0.01073093 -0.01042483]
   [-0.01327885 -0.01474966 -0.00524142]
   [ 0.01266256  0.00895064 -0.00601594]]]
 <NDArray 1x3x3 @gpu(1)>)

多GPU训练模型

我们先定义交叉熵损失函数。

In [6]:
loss = gloss.SoftmaxCrossEntropyLoss()

当我们使用多个GPU来训练模型时,gluon.Trainer会自动做数据并行,例如划分小批量数据样本并复制到各个GPU上,对各个GPU上的梯度求和再广播到所有GPU上。这样,我们就可以很方便地实现训练函数了。

In [7]:
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)
    net.initialize(init=init.Normal(sigma=0.01), ctx=ctx, force_reinit=True)
    trainer = gluon.Trainer(
        net.collect_params(), 'sgd', {'learning_rate': lr})
    for epoch in range(1, 6):
        start = time()
        for X, y in train_iter:
            gpu_Xs = gutils.split_and_load(X, ctx)
            gpu_ys = gutils.split_and_load(y, ctx)
            with autograd.record():
                ls = [loss(net(gpu_X), gpu_y) for gpu_X, gpu_y in zip(
                    gpu_Xs, gpu_ys)]
            for l in ls:
                l.backward()
            trainer.step(batch_size)
        nd.waitall()
        print('epoch %d, training time: %.1f sec'%(epoch, time() - start))
        test_acc = gb.evaluate_accuracy(test_iter, net, ctx[0])
        print('validation accuracy: %.4f'%(test_acc))

我们在2个GPU上训练模型。

In [8]:
train(num_gpus=2, batch_size=512, lr=0.3)
running on: [gpu(0), gpu(1)]
epoch 1, training time: 7.8 sec
validation accuracy: 0.7548
epoch 2, training time: 7.0 sec
validation accuracy: 0.6196
epoch 3, training time: 7.0 sec
validation accuracy: 0.8450
epoch 4, training time: 7.0 sec
validation accuracy: 0.8605
epoch 5, training time: 7.0 sec
validation accuracy: 0.8872

小结

  • 在Gluon中,我们可以很方便地进行多GPU计算,例如在多GPU上初始化模型参数和训练模型。

练习

  • 本节使用了ResNet-18。试试不同的迭代周期、批量大小和学习率。如果条件允许,使用更多GPU计算。
  • 有时候,不同的CPU/GPU的计算能力不一样,例如同时使用CPU和GPU,或者GPU之间型号不一样。这时候应该怎么办?

扫码直达讨论区