人工智能-模型训练之多GPU

每个新模型的并行计算都从零开始实现是无趣的。此外，优化同步工具以获得高性能也是有好处的。下面我们将展示如何使用深度学习框架的高级API来实现这一点。代码至少需要两个GPU来运行。

from mxnet import autograd, gluon, init, np, npx
from mxnet.gluon import nn
from d2l import mxnet as d2l

npx.set_np()

简单网络?

让我们使用一个比LeNet更有意义的网络，它依然能够容易地和快速地训练。我们选择的是（He et al., 2016）中的ResNet-18。因为输入的图像很小，所以稍微修改了一下。区别在于，我们在开始时使用了更小的卷积核、步长和填充，而且删除了最大汇聚层。

#@save
def resnet18(num_classes):
    """稍加修改的ResNet-18模型"""
    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(d2l.Residual(
                    num_channels, use_1x1conv=True, strides=2))
            else:
                blk.add(d2l.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