如何使用超过 1 个 GPU 训练 PyTorch 迁移学习教程

Question

我目前正在关注 PyTorch 迁移学习教程：https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html

我已经能够在 CPU 和 1 GPU 上完成教程和训练。

我正在使用 Google Cloud Platform Notebook Instances 并使用 4 个 NVIDIA Tesla k80 x 4 GPU。正是在这里，当我在 1 个以上的 GPU 上训练网络时，我遇到了服务器连接错误(无效响应：504(错误

model_ft = models.resnet18(pretrained=True)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, 2)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
## Using 4 GPUs
if torch.cuda.device_count() > 1:
model_ft = nn.DataParallel(model_ft)
model_ft = model_ft.to(device)
criterion = nn.CrossEntropyLoss()
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)
model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=25)

这个想法是利用数据并行函数来利用所有可用的GPU(即4个(来训练网络。

我在实施中遗漏了什么，请指教。

谢谢

Answer 1

IMO，使用Horovod进行多GPU训练是最容易的。以下是使用 Horovod 的 GPU 分布式训练脚本示例： https://github.com/horovod/horovod/blob/master/examples/pytorch_mnist.py

您需要安装OpenMPI(可能已经在包装盒上(，并且您需要在 python 环境中安装 Horovod(pip install horovod- 完整的安装说明在这里 https://github.com/horovod/horovod#install(。

然后，您将从horovodrun -np 4 python pytorch_mnist.py开始您的工作(这里有一些有关如何启动Horovod运行的文档：https://horovod.readthedocs.io/en/latest/mpirun.html(

这样，您不仅可以在具有多个 GPU 的一个节点上进行训练，还可以跨多个节点进行训练(例如，跨 2 个节点，每个节点有 4 个 GPU(。

使用 Horovod 进行分布式训练的要点是：

• Horovod 将启动您指示的任意数量的进程，因此在您的情况下为 4。每个进程将运行相同的脚本，仅在Horovod/MPI等级上有所不同。然后使用该等级来获取相应的 cuda 设备：

# Horovod: pin GPU to local rank.
torch.cuda.set_device(hvd.local_rank())
torch.cuda.manual_seed(args.seed)

• DistributedSampler用于在不同节点之间划分数据。hvd.rank()用于确保每个进程使用不同的数据分区，hvd.size()捕获总共有多少个进程。

train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset, num_replicas=hvd.size(), rank=hvd.rank())

• 用一个DistributedOptimizer包装你的优化器——它将负责在每个小批量结束时聚合各个进程的梯度：

# Horovod: wrap optimizer with DistributedOptimizer.
optimizer = hvd.DistributedOptimizer(optimizer,
named_parameters=model.named_parameters(),
compression=compression)

• 最重要的是，您不必修改实际模型和训练循环，这非常整洁。

样本中还有一些更有趣的事情(例如，随着数量或过程的增加而增加学习率，在开始时广播您的参数(。