如何使用kaggle中的两个GPU在pytorch中进行训练？

Question

我正在 Kaggle GPU 中训练模型。 但正如我所看到的，只有一个 GPU 正在工作。 我用普通的方法进行训练就像

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model = model.to(device)

如何同时使用两个 GPU？

Answer 1

使用多个 GPU 特定于机器学习库。我在 Pytorch 中进行图像分割时偶然发现了同样的问题。解决方案是将模块 torch.nn.DataParallel() 与模型一起使用。给定的代码可以更改如下：

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model = torch.nn.DataParallel(model, device_ids = [0,1]).to(device)

这里，

device_ids

是GPU的索引。假设如果您有 4 个 GPU，那么它将是

device_ids = [0,1,2,3]

或任何可能的索引。

结果 来了！.

PS：这是我在著名的 Stack Overflow 上的第一篇文章，请分享您的评论和观点。