我们有时候会遇到一种情况,即服务器上某节点包含多个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
以上就是我的解决过程,希望可以帮助到您