多 GPU 计算的 Gluon 实现

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

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

In [1]:
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]:
def resnet18(num_classes):  # 本函数已保存在 gluonbook 包中方便以后使用。
    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.48149228e-06 -8.33711965e-07 -1.63167783e-06 -6.36740253e-07
   -3.82161534e-06 -2.35140305e-06 -2.54696215e-06 -9.47841841e-08
   -6.90337401e-07  2.57562260e-06]
  [ 5.47108766e-06 -9.42465022e-07 -1.04940762e-06  9.80825234e-08
   -3.32518221e-06 -2.48629249e-06 -3.36428093e-06  1.04558694e-07
   -6.10015377e-07  2.03278250e-06]]
 <NDArray 2x10 @gpu(0)>,
 [[ 5.6176355e-06 -1.2837620e-06 -1.4605525e-06  1.8303012e-07
   -3.5511682e-06 -2.4371031e-06 -3.5731807e-06 -3.0974718e-07
   -1.1016568e-06  1.8909868e-06]
  [ 5.1418738e-06 -1.3729893e-06 -1.1520116e-06  1.1507541e-07
   -3.7372802e-06 -2.8289728e-06 -3.6477181e-06  1.5781598e-07
   -6.0733032e-07  1.9712006e-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 RuntimeError:
    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 来训练模型时,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: 15.5 sec, test_acc 0.89
epoch 2, training time: 14.7 sec, test_acc 0.91
epoch 3, training time: 14.4 sec, test_acc 0.92
epoch 4, training time: 14.8 sec, test_acc 0.93

然后尝试在 2 个 GPU 上训练。与上一节使用的 LeNet 相比,ResNet-18 的计算更加复杂,通讯时间与计算时间相比更短,因此 ResNet-18 的并行计算所获得的性能提升更佳。

In [8]:
train(num_gpus=2, batch_size=512, lr=0.2)
running on: [gpu(0), gpu(1)]
epoch 1, training time: 9.2 sec, test_acc 0.77
epoch 2, training time: 7.9 sec, test_acc 0.85
epoch 3, training time: 7.9 sec, test_acc 0.88
epoch 4, training time: 7.9 sec, test_acc 0.89

小结

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

练习

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

扫码直达讨论区