指定程序在哪个GPU上运行

摘要： 当本地（或服务器）有个多个GPU时，需要指定程序在指定GPU上运行，需要做以下设置。

一、在终端上指定GPU

在终端运行程序时，通过命令行指定：

CUDA_VISIBLE_DEVICES=1 python main.py
1

其中：1表示序号为1的GPU（注意：GPU的序号通常以0开始计数。）

二、在程序中指定GPU

在程序的开头，加入以下两行代码，指定运行的GPU，如程序main.py：

import os 
os.environ[“CUDA_VISIBLE_DEVICES”] = “1”
1
2

CUDA应用运行时，CUDA将遍历当前可见的设备，并从零开始为可见设备编号。可以同时指定多个设备（设备是真实存在的，不要指定不存在的设备）

import os 
os.environ[“CUDA_VISIBLE_DEVICES”] = “0, 1, 2, 3”
1
2

三、系统变量指定GPU

• Linux： export CUDA_VISIBLE_DEVICES=0
• windows: set CUDA_VISIBLE_DEVICES=0

四、pytorch中指定GPU

在程序中：

import torch
torch.cuda.set_device(1)
1
2

不过官方建议使用CUDA_VISIBLE_DEVICES，不建议使用 set_device 函数。

另外，再介绍几个pytorch关于cuda的函数：

# cuda是否可用，可用返回true, 不可用返回false
torch.cuda.is_available()

# 查看gpu可用的数量
torch.cuda.device_count()

# 返回gpu名称, 默认从0开始
torch.cuda.get_device_name(0) # 如rtx 3090

# 返回当前gpu的索引
torch.cuda.current_device()
1
2
3
4
5
6
7
8
9
10
11