ray.train.torch.get_devices

ray.train.torch.get_devices() → List[torch.device]

获取为此进程配置的正确的 torch 设备列表。

假设 CUDA_VISIBLE_DEVICES 已设置，并且是 ray.get_gpu_ids() 的超集。

示例

示例：在当前节点上启动了2个工作线程，每个工作线程使用1个GPU

os.environ["CUDA_VISIBLE_DEVICES"] == "2,3"
ray.get_gpu_ids() == [2]
torch.cuda.is_available() == True
get_devices() == [torch.device("cuda:0")]

示例：在当前节点上启动了4个工作进程，每个进程使用1个GPU

os.environ["CUDA_VISIBLE_DEVICES"] == "0,1,2,3"
ray.get_gpu_ids() == [2]
torch.cuda.is_available() == True
get_devices() == [torch.device("cuda:2")]

示例：在当前节点上启动了2个工作进程，每个工作进程使用2个GPU

os.environ["CUDA_VISIBLE_DEVICES"] == "0,1,2,3"
ray.get_gpu_ids() == [2,3]
torch.cuda.is_available() == True
get_devices() == [torch.device("cuda:2"), torch.device("cuda:3")]

PublicAPI (测试版): 此API目前处于测试阶段，在成为稳定版本之前可能会发生变化。