pytorch指定多块GPU运行代码

一. 参考链接：

pytorch指定多块GPU运行代码

二. 指定可见GPU

终端跑程序指定GPU

CUDA_VISIBLE_DEVICES=0,1,2,3 python XXX.py

或者在python文件中加入

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

多块GPU可选

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"

pytorch使用单GPU

model = model.cuda()

pytorch使用多GPU

model = torch.nn.DataParallel(model, device_ids=[0,1,2,3]).cuda()

使用了gpu0,1,2,3 共4块同时跑