Horovod 与 MXNet

Horovod 以类似的方式支持 Apache MXNet 和常规 TensorFlow。

查看完整的训练示例 MNIST 和 ImageNet。 以下脚本基于 Apache MXNet Gluon API 提供了一个简单的代码块框架。

import mxnet as mx
import horovod.mxnet as hvd
from mxnet import autograd

# Initialize Horovod
hvd.init()

# Pin GPU to be used to process local rank
context = mx.gpu(hvd.local_rank())
num_workers = hvd.size()

# Build model
model = ...
model.hybridize()

# Create optimizer
optimizer_params = ...
opt = mx.optimizer.create('sgd', **optimizer_params)

# Initialize parameters
model.initialize(initializer, ctx=context)

# Fetch and broadcast parameters
params = model.collect_params()
if params is not None:
    hvd.broadcast_parameters(params, root_rank=0)

# Create DistributedTrainer, a subclass of gluon.Trainer
trainer = hvd.DistributedTrainer(params, opt)

# Create loss function
loss_fn = ...

# Train model
for epoch in range(num_epoch):
    train_data.reset()
    for nbatch, batch in enumerate(train_data, start=1):
        data = batch.data[0].as_in_context(context)
        label = batch.label[0].as_in_context(context)
        with autograd.record():
            output = model(data.astype(dtype, copy=False))
            loss = loss_fn(output, label)
        loss.backward()
        trainer.step(batch_size)

注意

某些 MXNet 版本与 Horovod 不兼容：

• MXNet 1.4.0 及更早版本存在 GCC 兼容性问题。使用 MXNet 1.4.1 或更高版本配合 Horovod 0.16.2 或更高版本来避免这些兼容性问题。

• MXNet 1.5.1、1.6.0、1.7.0 和 1.7.0.post1 缺少 MKLDNN 头文件，因此它们无法与 Horovod 协同工作。请分别使用 1.5.1.post0、1.6.0.post0 和 1.7.0.post0 版本。

• MXNet 1.6.0.post0 和 1.7.0.post0 仅以 mxnet-cu101 和 mxnet-cu102 形式提供。