如何在pytorch中使用多个GPU？

Question

我正在学习pytorch并遵循此教程：https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html

我使用此命令来使用GPU。

device = torch.device("**cuda:0**" if torch.cuda.is_available() else "cpu")

但是，我想在 jupyter 中使用两个GPU，如下所示：

device = torch.device("**cuda:0,1**" if torch.cuda.is_available() else "cpu")

当然，这是错误的。 那么，我该怎么做？

Answer 1

使用多GPU就像将模型包装在DataParallel中并增加批处理大小一样简单。查看以下两个教程以快速入门：

• Multi-GPU Examples
• Data Parallelism

Answer 2

假设要在可用的GPU上分布数据（如果批处理大小为16个和2个GPU，则可能希望为每个GPU提供8个样本），而没有真正散布这些部分跨不同GPU的模型集。可以按照以下步骤进行操作：

如果要使用所有可用的GPU：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = CreateModel()

model= nn.DataParallel(model)
model.to(device)

如果要使用特定的GPU： （例如，使用4个GPU中的2个）

device = torch.device("cuda:1,3" if torch.cuda.is_available() else "cpu") ## specify the GPU id's, GPU id's start from 0.

model = CreateModel()

model= nn.DataParallel(model,device_ids = [1, 3])
model.to(device)

要通过设置OS环境变量来使用特定的GPU，

在执行程序之前，如下设置CUDA_VISIBLE_DEVICES变量：

export CUDA_VISIBLE_DEVICES=1,3（假设您要选择第二和第四GPU）

然后，在程序内，您就可以像使用所有GPU一样使用DataParallel()。 （类似于第一种情况）。可用于该程序的GPU在这里受OS环境变量的限制。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = CreateModel()

model= nn.DataParallel(model)
model.to(device)

在所有这些情况下，数据都必须映射到设备。

如果X和y是数据：

X.to(device)
y.to(device)

Answer 3

另一种选择是为PyTorch使用一些帮助程序库：

PyTorch Ignite库分布式GPU培训

在there中，有一个上下文管理器的概念，用于在以下位置进行分布式配置：

• nccl-在多个GPU上进行炬管本机分布式配置
• xla-tpu-TPU分布式配置

PyTorch Lightning多GPU训练

This是IMHO在CPU / GPU / TPU上进行培训的最佳选择，而无需更改原始PyTorch代码。

值得Catalyst来获得类似的分布式GPU选项。