【实用】多GPU未占满情况下，使用Python指定GPU运行代码

我们有时候会遇到一种情况，即服务器上某节点包含多个GPU，我们已经使用了一个或者几个GPU，还有部分GPU未使用，这个时候我们再次将代码上传服务器进行运行，发现不知道如何进行选择特定的GPU运行，且如果只通过默认运行，会发现报错，因为程序默认是从 GPU 0 开始的，如果GPU 0被占用了，则无法运行其他命令，会报 out of CUDA memory的异常。

遇到以上的问题，其实有两个解决方案，如下所示：

1 使用 nvidia-smi 进行查看GPU的信息，发现好多空载的GPU，这里注意每个GPU前面的编号 (0,1,2,…5)，这个就是后续指定GPU编号使用的。

1 在Python代码中直接声明

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3,4"#即指定使用 3 4 号GPU

import torch
#这个id 貌似只能是单个整数
torch.cuda.set_device(id)

2 在运行命令中选定GPU

在命令中使用 CUDA_VISIBLE_DEVICES=1,2 进行选定，如下，例子中为指定1，2两块GPU进行运行该程序

CUDA_VISIBLE_DEVICES=1,2 python -m torch.distributed.launch --nproc_per_node 2 --master_port 12345  main.py 

以上就是我的解决过程，希望可以帮助到您