需要知道的几个点:
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 卡进行后续推理。