相关阅读

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.