【深度学习】Python使用指定gpu运行代码

今天药忘吃喽~ 2024-02-20 10:03 66阅读 0赞

### 命令行指定显卡GPU运行python脚本 ###

在大型机构分配的服务器集群中，需要使用GPU的程序默认都会在第一张卡上进行，如果第一张卡倍别人占用或者显存不够的情况下，程序就会报错说没有显存容量，所以能够合理地利用GPU资源能帮助你更快更好地跑出实验效果。  
**1、指定使用GPU0运行脚本（默认是第一张显卡, 0代表第一张显卡的id,其他的以此类推）**  
第一种方式：

CUDA_VISIBLE_DEVICES=0 python ***.py

第二种方式：在python脚本文件中添加如下内容进行指定：

import os
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

如果在脚本文件中已经写明了第二种方式，则不支持第一种方式的指定。  
**2、指定使用多张显卡运行脚本**  
在GPU的id为0和1的两张显卡上运行\*\*\*.py程序：

CUDA_VISIBLE_DEVICES=0,1 python ***.py

**3、在单张显卡的情况下开启多个进程运行脚本**  
命令行程序如下：（其中NUM\_GPUS\_YOU\_HAVE代表运行的进程数）

说明：程序会根据卡的容量分配成两个程序，不一定在两张卡，有时候一张卡的容量够则会在该张卡上开启新的进程。

python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE ***.py

**4、查看GPU状态**  
1、单次查看

nvidia-smi

2、隔一秒查看GPU状态：

watch -n 1 nvidia-smi

### 使用指定gpu运行代码 ###

#### 一、前提 ####

1、在命令行使用nvidia-smi查看gpu设备情况，当存在空闲设备时才能用，否则会出现运行内存不够出错的情况（具体参考文章GPU之nvidia-smi命令详解）；  
2、安装好cuda和cudcnn（具体步骤请参考：①windows：CUDA安装教程（超详细）），②linux：linux安装CUDA+cuDNN）  
在命令行输入nvcc -V或者nvcc --version检查是否安装成功：

#### 二、命令行直接指定 ####

运行python文件时，在前面加上 CUDA\_VISIBLE\_DEVICES=xxx，xxx是你要指定的gpu序号。如CUDA\_VISIBLE\_DEVICES=0,2 python extract\_masks.py表示指定了0、2号两个gpu。注意，序号是从0开始排的，所以如果你只有一个gpu，那么就是CUDA\_VISIBLE\_DEVICES=0 python extract\_masks.py。

#### 三、在python文件中指定 ####

在需要运行的python文件最开始加上两行如下，等号后面的数字同上所述。

import os
    os.environ["CUDA_VISIBLE_DEVICES"] = "0,2,3"
    # 注意：这两行代码必须在文件的最开头，在加载各种包之前

#### 四、如何使用 ####

在python文件中，定义需要加速的模型之后，加上：

model = ......
    model.train(True) or model.train(False) # 看你是要训练还是测试
    model.to('cuda') # 或者model.cuda()
    # 后面需要输入model的变量也是需要.to('cuda')或者.cuda()的，不然会报错既用了cpu又用gpu，不兼容