Pytorch基础:环境变量CUDA_VISIBLE_DEVICES 相关阅读Pytorch基础https://blog.csdn.net/weixin_45791458/category_12457644.html?spm1001.2014.3001.5482CUDA_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_1torch.tensor([1,2,3,4,5], devicecpu) print(tensor_1) tensor_1tensor_1.to(cuda:0) print(tensor_1) tensor_1tensor_1.to(cuda:1) print(tensor_1) tensor_1tensor_1.to(cuda:2) print(tensor_1) tensor_2torch.tensor([11,22,33,44,55], devicecpu) print(tensor_2) tensor_2tensor_2.to(cuda:0) print(tensor_2) tensor_2tensor_2.to(cuda:1) print(tensor_2) tensor_2tensor_2.to(cuda:2) print(tensor_2) tensor_3torch.tensor([111,222,333,444,555], devicecpu) print(tensor_3) device_0torch.device(cuda:0) tensor_3tensor_3.to(device_0) print(tensor_3) device_1torch.device(cuda:1) tensor_3tensor_3.to(device_1) print(tensor_3) device_2torch.device(cuda:2) tensor_3tensor_3.to(device_2) print(tensor_3)输出 tensor([1, 2, 3, 4, 5]) tensor([1, 2, 3, 4, 5], devicecuda:0) tensor([1, 2, 3, 4, 5], devicecuda:1) tensor([1, 2, 3, 4, 5], devicecuda:2) tensor([11, 22, 33, 44, 55]) tensor([11, 22, 33, 44, 55], devicecuda:0) tensor([11, 22, 33, 44, 55], devicecuda:1) tensor([11, 22, 33, 44, 55], devicecuda:2) tensor([111, 222, 333, 444, 555]) tensor([111, 222, 333, 444, 555], devicecuda:0) tensor([111, 222, 333, 444, 555], devicecuda:1) tensor([111, 222, 333, 444, 555], devicecuda: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_DEVICES1,2 $ export CUDA_VISIBLE_DEVICES1,2 $ export CUDA_VISIBLE_DEVICES1,2 $ export # 查看是否成功设置环境变量 ****** declare -x CUDA_VISIBLE_DEVICES1,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 atorch.tensor([1, 2, 3, 4, 5], devicecuda: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_BLOCKING1. 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_DEVICES2,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设备数量 atorch.tensor([1, 2, 3, 4, 5], devicecuda: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_BLOCKING1. Compile with TORCH_USE_CUDA_DSA to enable device-side assertions.