在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