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

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

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

我正在使用Google Cloud Platform Notebook实例，并使用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）来训练网络。

我在实施中缺少什么，请指教。

谢谢

最佳答案

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的一个节点上进行训练，而且还可以在多个节点上进行训练（例如，在每个具有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)

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


样本中还有其他一些有趣的事情（例如，随着学习次数或过程的增加，学习率会提高，开始时会广播您的参数）。

关于python - 如何使用超过1个GPU训练PyTorch转移学习教程，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/57403213/