使用CUDA_VISIBLE_DEVICES设置显卡

一. 读卡规则

当服务器有多个GPU时,设置CUDA_VISIBLE_DEVICES 环境变量可以改变CUDA程序所能使用的GPU设备。在默认情况下,标号为0的显卡为主卡,如主机中有4块显卡,那么每张显卡的默认标号为[0,1,2,3]。多卡设置规则如下:

Environment Variable SyntaxResults
CUDA_VISIBLE_DEVICES=1Only device 1 will be seen
CUDA_VISIBLE_DEVICES=0,1Devices 0 and 1 will be visible
CUDA_VISIBLE_DEVICES="0,1"Same as above, quotation marks are optiona
CUDA_VISIBLE_DEVICES=0,2,3lDevices 0, 2, 3 will be visible;device 1 is masked

CUDA应用运行时,CUDA将遍历当前可见的设备,并从零开始为可见设备编号。第一种情况,卡1设置为主卡,但CUDA遍历时会设置为可见编号0。最后一种情况,设备0,2,3将显示为设备0,1,2。如果将字符串的顺序更改为“2,3,0”,则设备2,3,0将分别被设置为0,1,2。如果为CUDA_VISIBLE_DEVICES 设置了不存在的设备,所有实际设备将被隐藏,CUDA 应用将无法使用GPU设备;如果设备序列是存在和不存在设备的混合,那么不存在设备前的所有存在设备将被重新编号,不存在设备之后的所有设备将被屏蔽。当前可见的(重新编号后的)设备可使用CUDA 程序来查看。

举例:

import torch
print(torch.cuda.current_device())

运行:

结果:

虽然指定使用设备为卡1,但CUDA会编号为0,因此显示的编号还是0

二. 设置GPU

1. 临时设置

1.1. 直接设置

export CUDA_VISIBLE_DEVICES=1

1.2. 运行bash脚本时指定GPU

CUDA_VISIBLE_DEVICES=gpu_ids sh test.sh

1.3. 在代码中指定GPU

import os 
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_ids 

1.4. 运行程序时指定GPU

CUDA_VISIBLE_DEVICES=gpu_ids python test.py

注意:这种设置方法一定要在第一次使用 cuda 之前进行设置

2. 永久设置

管理员身份运行

~/.bashrc

在最后加上

export CUDA_VISIBLE_DEVICES=gpu_ids

然后

source ~/.bashrc