pytorch - 具有在 PyTorch 中指定的 gpu 设备 ID 的 DistributedDataParallel

Question

我想在具有 8 个 GPU 的单台机器上通过 DistributedDataParallel 训练我的模型。但我想在设备 ID 为 4、5、6、7 的四个指定 GPU 上训练我的模型。

如何为 DistributedDataParallel 指定 GPU 设备 ID？

我认为这种情况下的世界大小将是 4，但这种情况下的排名应该是多少？

Answer 1

您可以设置环境变量CUDA_VISIBLE_DEVICES。Torch 将读取此变量并仅使用其中指定的 GPU。您可以直接在您的 python 代码中执行此操作，如下所示：

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '4, 5, 6, 7'

在以任何方式初始化torch之前请注意执行此命令，否则该语句将不会生效。另一个选项是在 shell 中启动脚本之前临时设置环境变量：

CUDA_VISIBLE_DEVICES=4,5,6,7 python your_script.py