pytorch指定多块GPU运行代码_os.environ["cuda_visible_devices"] = "0

查看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块同时跑

其他

训练后保存的模型参数，再次导入模型时，出现报错：

Missing key(s) in state_dict: ...
1

解决办法：
在导入参数前加上
import torch.nn as nn
from torch.backends import cudnn

model = torch.nn.DataParallel(model)
cudnn.benchmark = True
1
2