原文链接《pytorch指定多块GPU运行代码》, 一起来围观吧 /OneQuestionADay/article/details/111691486?utm_source=app
查看GPU状态和信息
nvidia-smi
1
在这里插入图片描述在这里插入图片描述
定时刷新状态,-n(秒),每秒刷新一次
watch -n 1 nvidia-smi
1
指定可见GPU
终端跑程序指定GPU
CUDA_VISIBLE_DEVICES=0 python XXX.py
1
或者在python文件中加入
import os
os.environ[“CUDA_VISIBLE_DEVICES”] = “0”
1
2
多块GPU可选
os.environ[“CUDA_VISIBLE_DEVICES”] = “0,1,2,3”
1
pytorch使用单GPU
model = model.cuda()
1
pytorch使用多GPU
model = torch.nn.DataParallel(model, device_ids=[0,1,2,3]).cuda()
1
使用了gpu0,1,2,3 共4块同时跑