多 GPU 计算的 Gluon 实现

在 Gluon 中,我们可以很方便地使用数据并行进行多 GPU 计算。比方说,我们并不需要自己实现“多 GPU 计算”一节里介绍的多 GPU 之间同步数据的辅助函数。先导入本节实验需要的包或模块。同上一节,运行本节中的程序需要至少两块 GPU。

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

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
import time

多 GPU 上初始化模型参数

我们使用 ResNet-18 来作为本节的样例模型。由于本节的输入图像使用原尺寸(未放大),这里的模型构造与“残差网络(ResNet)”一节中的 ResNet-18 构造稍有不同。这里的模型在一开始使用了较小的卷积核、步幅和填充,并去掉了最大池化层。

In [2]:
# 本函数已保存在 gluonbook 包中方便以后使用。
def resnet18(num_classes):
    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 = nn.Sequential()
    # 这里使用了较小的卷积核、步幅和填充,并去掉了最大池化层。
    net.add(nn.Conv2D(64, kernel_size=3, strides=1, padding=1),
            nn.BatchNorm(), nn.Activation('relu'))
    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.4814936e-06 -8.3371094e-07 -1.6316770e-06 -6.3674099e-07
   -3.8216162e-06 -2.3514044e-06 -2.5469599e-06 -9.4784696e-08
   -6.9033558e-07  2.5756231e-06]
  [ 5.4710872e-06 -9.4246496e-07 -1.0494070e-06  9.8081841e-08
   -3.3251815e-06 -2.4862918e-06 -3.3642798e-06  1.0455864e-07
   -6.1001344e-07  2.0327841e-06]]
 <NDArray 2x10 @gpu(0)>,
 [[ 5.6176345e-06 -1.2837586e-06 -1.4605541e-06  1.8302967e-07
   -3.5511653e-06 -2.4371013e-06 -3.5731798e-06 -3.0974860e-07
   -1.1016571e-06  1.8909889e-06]
  [ 5.1418697e-06 -1.3729932e-06 -1.1520088e-06  1.1507450e-07
   -3.7372811e-06 -2.8289724e-06 -3.6477197e-06  1.5781629e-07
   -6.0733043e-07  1.9712013e-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 训练模型

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

In [6]:
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})
    loss = gloss.SoftmaxCrossEntropyLoss()
    for epoch in range(4):
        start = time.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()
        train_time = time.time() - start
        test_acc = gb.evaluate_accuracy(test_iter, net, ctx[0])
        print('epoch %d, training time: %.1f sec, test_acc %.2f' % (
            epoch + 1, train_time, test_acc))

首先在单 GPU 上训练。

In [7]:
train(num_gpus=1, batch_size=256, lr=0.1)
running on: [gpu(0)]
epoch 1, training time: 178.0 sec, test_acc 0.90
epoch 2, training time: 170.7 sec, test_acc 0.90
epoch 3, training time: 76.4 sec, test_acc 0.92
epoch 4, training time: 63.6 sec, test_acc 0.93

然后尝试 2 个 GPU。比上一节使用的 LeNet,ResNet-18 计算更加复杂,其并行效果更佳。

In [8]:
train(num_gpus=2, batch_size=512, lr=0.2)
running on: [gpu(0), gpu(1)]
epoch 1, training time: 33.0 sec, test_acc 0.80
epoch 2, training time: 32.3 sec, test_acc 0.88
epoch 3, training time: 32.1 sec, test_acc 0.91
epoch 4, training time: 32.3 sec, test_acc 0.91

小结

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

练习

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

扫码直达讨论区