多 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, nn, utils as gutils
from time import time

多 GPU 上初始化模型参数

我们使用 ResNet-18 来作为本节的样例模型。我们将 resnet18 函数定义在 gluonbook 包中供后面章节调用。

In [2]:
def resnet18(num_classes):
    net = nn.Sequential()
    net.add(nn.Conv2D(64, kernel_size=3, strides=1, padding=1),
            nn.BatchNorm(), nn.Activation('relu'),
            nn.MaxPool2D(pool_size=3, strides=2, padding=1))

    def resnet_block(num_channels, num_residuals, first_block=False):
        blk = nn.Sequential()
        for i in range(num_residuals):
            if i == 0 and not first_block:
                blk.add(gb.Residual(num_channels, use_1x1conv=True,
                                    strides=2))
            else:
                blk.add(gb.Residual(num_channels))
        return blk

    net.add(resnet_block(64, 2, first_block=True),
            resnet_block(128, 2),
            resnet_block(256, 2),
            resnet_block(512, 2))
    net.add(nn.GlobalAvgPool2D(), nn.Dense(num_classes))
    return net

net = 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]:
(
 [[  5.88889543e-06  -3.19231913e-06  -1.56602857e-06   2.42315537e-07
    -5.51294079e-06  -2.40595455e-06  -2.95329119e-06   6.47307672e-07
    -3.02829903e-07   2.56456224e-06]
  [  5.85794623e-06  -3.91601498e-06  -1.02782406e-06  -3.87587988e-07
    -5.79865173e-06  -3.17184140e-06  -2.32055481e-06   1.00785905e-06
    -8.17856517e-07   2.59511353e-06]]
 <NDArray 2x10 @gpu(0)>,
 [[  5.54097005e-06  -3.90407240e-06  -1.33178651e-06  -6.91773039e-07
    -5.98816723e-06  -2.48047218e-06  -2.46669651e-06  -3.57805902e-07
    -5.74706633e-07   3.10877112e-06]
  [  5.53567588e-06  -3.93424261e-06  -8.16369777e-07  -3.07119990e-07
    -5.28076680e-06  -3.08374274e-06  -2.77396771e-06   8.90484444e-07
    -6.43086764e-07   2.54331326e-06]]
 <NDArray 2x10 @gpu(1)>)

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

In [5]:
weight = net[0].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: 14.6 sec
validation accuracy: 0.7089
epoch 2, training time: 13.0 sec
validation accuracy: 0.7943
epoch 3, training time: 12.8 sec
validation accuracy: 0.8301
epoch 4, training time: 12.8 sec
validation accuracy: 0.7657
epoch 5, training time: 12.9 sec
validation accuracy: 0.8772

小结

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

练习

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

扫码直达 讨论区