如何在 pytorch 中使用多個 GPU？

Question

我用這個命令來使用一個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包裝模型並增加批量大小一樣簡單。 查看這兩個教程以快速入門：

• 多 GPU 示例
• 數據並行

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)

要通過設置操作系統環境變量來使用特定 GPU：

在執行程序之前，設置CUDA_VISIBLE_DEVICES變量如下：

export CUDA_VISIBLE_DEVICES=1,3 （假設您要選擇第 2 個和第 4 個 GPU）

然后，在程序中，您可以像使用所有 GPU 一樣使用DataParallel() 。 （類似於第一種情況）。 這里可供程序使用的 GPU 受操作系統環境變量的限制。

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 訓練

其中有一個用於分布式配置的上下文管理器的概念：

• nccl - 多個 GPU 上的火炬本機分布式配置
• xla-tpu - TPU 分布式配置

PyTorch Lightning Multi-GPU 訓練

恕我直言， 這可能是在不更改原始 PyTorch 代碼的情況下在 CPU/GPU/TPU 上訓練的最佳選擇。

對於類似的分布式 GPU 選項，值得一試Catalyst 。

Answer 4

2022 年，PyTorch 說：

建議使用 DistributedDataParallel，而不是這個 class，來做多 GPU 訓練，即使只有一個節點。 請參閱：使用 nn.parallel.DistributedDataParallel 而不是多處理或 nn.DataParallel 和分布式數據並行。

在https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html#torch.nn.DataParallel

因此，似乎我們應該使用DistributedDataParallel ，而不是DataParallel 。

Answer 5

當我運行 naiveinception_googlenet 時，上述方法對我不起作用。 以下方法解決了我的問題。

導入操作系統

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"

os.environ["CUDA_VISIBLE_DEVICES"]="0,3" # 指定要使用的 GPU

Answer 6

如果您只想在特定 GPU 上運行您的代碼（例如，僅在 GPU id 2 和 3 上），那么您可以在從終端觸發 python 代碼時指定使用 CUDA_VISIBLE_DEVICES=2,3 變量。

CUDA_VISIBLE_DEVICES=2,3 python lstm_demo_example.py --epochs=30 --lr=0.001

在代碼中，將其保留為：

device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')
model = LSTMModel()
model = nn.DataParallel(model)
model = model.to(device)

資料來源： https ://glassboxmedicine.com/2020/03/04/multi-gpu-training-in-pytorch-data-and-model-parallelism/