【实用】多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 

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

猜你喜欢

转载自blog.csdn.net/qq_29750461/article/details/128985985