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

最新推荐文章于 2024-09-23 15:19:09 发布

原创最新推荐文章于 2024-09-23 15:19:09 发布 · 1.4w 阅读

39 ·

CC 4.0 BY-SA版权

PyTorch 基础专栏收录该内容

10 篇文章

订阅专栏

本文探讨了在使用PyTorch时，GPU编号与nvidia-smi显示不一致的问题，重点在于如何通过设置CUDA_DEVICE_ORDER为PCI_BUS_ID来统一设备标识。提供了Windows和Ubuntu下的解决方案，适用于TensorFlow和NVIDIA显卡管理。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

问题描述

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