多 GPU 计算的简洁实现¶
在 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.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 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: 63.7 sec, test_acc 0.88
epoch 2, training time: 59.9 sec, test_acc 0.90
epoch 3, training time: 59.9 sec, test_acc 0.91
epoch 4, training time: 59.9 sec, test_acc 0.92
然后尝试在 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: 31.7 sec, test_acc 0.78
epoch 2, training time: 30.6 sec, test_acc 0.86
epoch 3, training time: 30.4 sec, test_acc 0.89
epoch 4, training time: 30.7 sec, test_acc 0.88
小结¶
- 在 Gluon 中，我们可以很方便地进行多 GPU 计算，例如在多 GPU 上初始化模型参数和训练模型。
练习¶
- 本节使用了 ResNet-18。试试不同的迭代周期、批量大小和学习率。如果条件允许，使用更多 GPU 计算。
- 有时候，不同设备的计算能力不一样，例如同时使用 CPU 和 GPU，或者不同 GPU 之间型号不一样。这时候应该如何划分小批量到不同的 CPU 或 GPU？
