PyTorch 代码中 GPU 编号与 nvidia-smi 命令中的 GPU 编号不一致问题解决方法

问题描述

• nvidia-smi -L 和 PyTorch 代码中 GPU 编号不一致

• nvidia-smi -L 查看可用 GPU 列表

GPU 0: Quadro Kxxx (UUID: GPU-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx)
GPU 1: GeForce RTX xxxx (UUID: GPU-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx)

• PyTorch 代码中查看可用 GPU 列表

device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
print("GPU 编号: {}".format(device))
print("GPU 名称: {}".format(torch.cuda.get_device_name(1)))


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("GPU 编号: {}".format(device))
print("GPU 名称: {}".format(torch.cuda.get_device_name(0)))

输出：
GPU 编号: cuda:1
GPU 名称: Quadro Kxxx
GPU 编号: cuda:0
GPU 名称: GeForce RTX xxxx

• 可以看出 nvidia-smi 命令中的 GPU 编号与 PyTorch 代码中的 GPU 编号正好相反（两块显卡的情况）

原因

• nvidia-smi 下的 GPU 编号默认使用 PCI_BUS_ID，而 PyTorch 代码默认情况下设备排序是 FASTEST_FIRST

解决方法

• 要使 PyTorch 代码中的 GPU 编号与 nvidia_smi 保持一致，需要讲 expert CUDA_DEVICE_ORDER=PCI_BUS_ID 添加到bashrc

• Windows 下可以在 PyTorch 代码中添加如下：

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"

• Ubuntu 下使用，可以直接添加如下，也可以 CUDA_DEVICE_ORDER=“PCI_BUS_ID” 添加到 bashrc 文件中

export CUDA_DEVICE_ORDER="PCI_BUS_ID"

PS：TensorFlow 如遇相同问题也可以参照修改

参考

• https://stackoverflow.com/questions/26123252/inconsistency-of-ids-between-nvidia-smi-l-and-cudevicegetname
• https://github.com/matterport/Mask_RCNN/issues/109