关于 torch 的 device id 与真实 GPU id 的关系

需要知道的几个点：

• cuda:{id} 中的 id 并不一定是真实硬件的GPU id，而是运行时可用的 GPU id（从0开始计数）
• torch.cuda.device_count() 可查看运行时可用的 GPU 数量
• torch.cuda.get_device_name(i) 可获取第 i 个 device 的 name

测试代码。server.py：

device_count = torch.cuda.device_count()

for i in range(device_count):
    print(f"Device {
      
      i}: {
      
      torch.cuda.get_device_name(i)}")

device = torch.device(f"cuda:1" if torch.cuda.is_available() else "cpu")

比如，执行命令：

CUDA_VISIBLE_DEVICES=4,5 python server.py

输出结果：

Device 0: NVIDIA GeForce RTX 3090
Device 1: NVIDIA GeForce RTX 3090

此时，代码会认为可见的第一个GPU（即 cuda:0）是4卡，可见的第二个GPU（即 cuda:1）是5卡。server.py 代码中的 device 就是用了 5 卡进行后续推理。