环境变量一般存在于操作系统内部，应用程序启动之后，它会从操作系统调用其存储的某些环境变量，从而完成某种控制行为。
CUDA_VISIBLE_DEVICES环境变量是专门提供给CUDA程序使用的环境变量，CUDA是NVIDIA开发的并行计算平台和编程模型，我们平时安装的PyTorch绝大多数都是基于CUDA进行模型训练的。
CUDA_VISIBLE_DEVICES环境变量用于指定CUDA应用程序将在哪些GPU设备上运行，通常用于多GPU系统，对于单GPU系统则没有意义。通过设置CUDA_VISIBLE_DEVICES环境变量，可以限制应用程序访问的GPU 设备，以便在多任务或多用户环境中更好地管理和分配GPU资源。
CUDA_VISIBLE_DEVICES 的值是一个以英文逗号分隔的 GPU 设备索引表，例如：0,1,2。这表示应用程序将只能在索引为 0、1、2 的 GPU 设备上运行，而忽略其他 GPU 设备。如果用户没有显式设置 CUDA_VISIBLE_DEVICES 的值，应用程序将默认使用所有可用的 GPU 设备。
有时候，我们只想使用其中的一个GPU进行计算。我们可以通过设置CUDA_VISIBLE_DEVICES环境变量来指定使用的GPU编号。
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
在上述例子中，我们将使用编号为0的GPU进行计算。需要注意的是，当你设置了：
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
你实际上是告诉CUDA只看到一个GPU，即设备ID为3的GPU。因此，对于你的PyTorch程序来说，此GPU对应的cuda设备索引为0。
CUDA_VISIBLE_DEVICES环境变量有两大作用：
决定GPU在CUDA程序中是否可见
影响GPU设备名映射cuda设备编号
我们知道使用nvidia-smi命令可以查询本机GPU的相关信息，如下所示：
C:\Users\488xx66>nvidia-smi
Wed Oct 2 20:36:52 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 NVIDIA .. RTX 4090 Off | 00000000:02:00.0 Off | 0 |
| N/A 44C P0 79W / 250W | 12178MiB / 16384MiB | 42% Default |
| | | N/A |
+-----------------------------------------+----------------------+----------------------+
| 1 NVIDIA .. RTX 4090 Off | 00000000:03:00.0 Off | 0 |
| N/A 49C P0 81W / 250W | 12214MiB / 16384MiB | 63% Default |
| | | N/A |
+-----------------------------------------+----------------------+----------------------+
| 2 NVIDIA .. RTX 4090 Off | 00000000:82:00.0 Off | 0 |
| N/A 35C P0 31W / 250W | 2932MiB / 16384MiB | 0% Default |
| | | N/A |
+-----------------------------------------+----------------------+----------------------+
可以看到我的主机有三个GPU，分别被编号为0、1、2。下面使用torch.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型号: NVIDIA GeForce RTX 4090
CUDA设备1型号: NVIDIA GeForce RTX 4090
CUDA设备2型号: NVIDIA GeForce RTX 4090
下面给大家演示一下CUDA_VISIBLE_DEVICES环境变量的两大作用。如下所示，仅设置GPU1、GPU2可见，且将GPU1当做cuda:0设备，GPU2当做cuda:1设备：
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"
# 查看cuda设备数量
print("CUDA设备数量:", torch.cuda.device_count())
# 无法使用cuda2设备
a=torch.tensor([1, 2, 3, 4, 5], device="cuda:2")
上述代码设置GPU1和GPU2可见，且将GPU1当做cuda:0设备，GPU2当做cuda:1设备，所以此时无法使用cuda:2设备。由于环境变量CUDA_VISIBLE_DEVICES中编号的顺序会影响cuda设备的映射编号，我们甚至可以将GPU0当做cuda:2设备，GPU1当做cuda:1设备，GPU2当做cuda:0设备，如下所示。
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2,1,0"