本文主要是介绍Pytorch基础：环境变量CUDA_VISIBLE_DEVICES，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！
相关阅读
Pytorch基础
https://blog.csdn.net/weixin_45791458/category_12457644.html?spm=1001.2014.3001.5482
CUDA_VISIBLE_DEVICES这个环境变量可以影响CUDA能识别到的GPU，并影响它映射到的cuda设备编号。
首先我们知道使用nvidia-smi命令可以查询本机GPU的相关信息，如下所示。
$ nvidia-smi
Sun May 12 22:13:43 2024
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.161.07 Driver Version: 535.161.07 CUDA Version: 12.2 |
|-----------------------------------------+----------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+======================+======================|
| 0 Tesla P100-PCIE-16GB Off | 00000000:02:00.0 Off | 0 |
| N/A 44C P0 79W / 250W | 12178MiB / 16384MiB | 42% Default |
| | | N/A |
+-----------------------------------------+----------------------+----------------------+
| 1 Tesla P100-PCIE-16GB Off | 00000000:03:00.0 Off | 0 |
| N/A 49C P0 81W / 250W | 12214MiB / 16384MiB | 63% Default |
| | | N/A |
+-----------------------------------------+----------------------+----------------------+
| 2 Tesla P100-PCIE-16GB Off | 00000000:82:00.0 Off | 0 |
| N/A 35C P0 31W / 250W | 2932MiB / 16384MiB | 0% Default |
| | | N/A |
+-----------------------------------------+----------------------+----------------------+可以看到我的主机有三个GPU，分别被编号为0、1、2。在下面的python代码中，我们使用pytorch.cuda模块下提供的函数，查询了本机拥有的cuda设备数量，并根据编号显示了它们的型号。
import torch# 检查是否有可用的CUDA设备
if torch.cuda.is_available():# 获取设备数量device_count = torch.cuda.device_count()print("CUDA设备数量:", device_count)# 遍历每个设备并输出型号for i in range(device_count):print("CUDA设备{}型号:".format(i), torch.cuda.get_device_name(i))
else:print("CUDA不可用")输出：
CUDA设备数量: 3
CUDA设备0型号: Tesla P100-PCIE-16GB
CUDA设备1型号: Tesla P100-PCIE-16GB
CUDA设备2型号: Tesla P100-PCIE-16GB上面的cuda设备编号，就对应着nvidia-smi命令所查询到的GPU编号，下面所示的几种方式可以将张量在指定编号的cuda设备之间转移，其中cuda:0、cuda:1、cuda:2中的数字指的就是cuda设备的编号。
import torchtensor_1=torch.tensor([1,2,3,4,5], device='cpu')
print(tensor_1)tensor_1=tensor_1.to('cuda:0')
print(tensor_1)
tensor_1=tensor_1.to('cuda:1')
print(tensor_1)
tensor_1=tensor_1.to('cuda:2')
print(tensor_1)tensor_2=torch.tensor([11,22,33,44,55], device='cpu')
print(tensor_2)tensor_2=tensor_2.to('cuda:0')
print(tensor_2)
tensor_2=tensor_2.to('cuda:1')
print(tensor_2)
tensor_2=tensor_2.to('cuda:2')
print(tensor_2)tensor_3=torch.tensor([111,222,333,444,555], device='cpu')
print(tensor_3)device_0=torch.device('cuda:0')
tensor_3=tensor_3.to(device_0)
print(tensor_3)
device_1=torch.device('cuda:1')
tensor_3=tensor_3.to(device_1)
print(tensor_3)
device_2=torch.device('cuda:2')
tensor_3=tensor_3.to(device_2)
print(tensor_3)输出：
tensor([1, 2, 3, 4, 5])
tensor([1, 2, 3, 4, 5], device='cuda:0')
tensor([1, 2, 3, 4, 5], device='cuda:1')
tensor([1, 2, 3, 4, 5], device='cuda:2')
tensor([11, 22, 33, 44, 55])
tensor([11, 22, 33, 44, 55], device='cuda:0')
tensor([11, 22, 33, 44, 55], device='cuda:1')
tensor([11, 22, 33, 44, 55], device='cuda:2')
tensor([111, 222, 333, 444, 555])
tensor([111, 222, 333, 444, 555], device='cuda:0')
tensor([111, 222, 333, 444, 555], device='cuda:1')
tensor([111, 222, 333, 444, 555], device='cuda:2')有了上面的基础，下面来谈谈环境变量CUDA_VISIBLE_DEVICES，在一般情况下，它是没有定义的，可以使用下面的Python代码进行检测。
import os
print(os.environ["CUDA_VISIBLE_DEVICES"])
输出：
KeyError: 'CUDA_VISIBLE_DEVICES'可以在Bash中使用export命令，定义环境变量CUDA_VISIBLE_DEVICES（设置变量为导出状态），有关export命令的内容，可以查看Linux：导出环境变量命令export。这种情况下，在该Bash下运行的所有python子进程，都会继承这个环境变量，如下所示。
# 以下三种方式均可以成功设置
# 仅设置GPU1、GPU2可见，且将GPU1当做cuda0设备，GPU2当做cuda1设备
$ export CUDA_VISIBLE_DEVICES='1,2'
$ export CUDA_VISIBLE_DEVICES="1,2"
$ export CUDA_VISIBLE_DEVICES=1,2$ export # 查看是否成功设置环境变量
******
declare -x CUDA_VISIBLE_DEVICES="1,2"
******$ python
Python 3.9.18 (main, Sep 11 2023, 13:41:44)
[GCC 11.2.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import os
>>> print(os.environ["CUDA_VISIBLE_DEVICES"])
1,2
>>> import torch
>>> print("CUDA设备数量:", torch.cuda.device_count()) # 查看cuda设备数量
>>> CUDA设备数量: 2
>>> a=torch.tensor([1, 2, 3, 4, 5], device='cuda:2') # 无法使用cuda2设备
RuntimeError: CUDA error: invalid device ordinal
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.上例展示了，如何仅设置GPU1和GPU2可见，且将GPU1当做cuda0设备，GPU2当做cuda1设备，所以此时无法使用cuda2设备。由于环境变量CUDA_VISIBLE_DEVICES中编号的顺序也会影响cuda设备的映射编号，我们甚至可以将GPU0当做cuda2设备，GPU1当做cuda1设备，GPU2当做cuda0设备，如下所示。
$ export CUDA_VISIBLE_DEVICES='2,1,0' 还可以在python内部使用os模块直接设置本python进程的环境变量，如下所示。
import os
import torch
os.environ["CUDA_VISIBLE_DEVICES"]='2' # 注意：这里必须使用字符串
print(os.environ["CUDA_VISIBLE_DEVICES"])
print("CUDA设备数量:", torch.cuda.device_count()) # 查看cuda设备数量
a=torch.tensor([1, 2, 3, 4, 5], device='cuda:1') # 无法使用cuda1设备输出：
2
CUDA设备数量: 1
RuntimeError: CUDA error: invalid device ordinal
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
这篇关于Pytorch基础：环境变量CUDA_VISIBLE_DEVICES的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！
