python 深度学习模型训练 多GPU下如何调用

原文链接《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块同时跑