相关阅读
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 torch tensor_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.)
