我想在一台拥有 8 个 GPU 的机器上通过 DistributedDataParallel 训练我的模型。但我想在设备 ID 为 4、5、6、7 的四个指定 GPU 上训练我的模型。
如何为DistributedDataParallel指定GPU设备ID？
我认为这种情况下的世界大小将为 4，但是这种情况下的排名应该是多少？
最佳答案
您可以设置环境变量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
关于pytorch - DistributedDataParallel 与 PyTorch 中指定的 GPU 设备 ID，我们在Stack Overflow上找到一个类似的问题： https://stackoverflow.com/questions/69703158/