查看GPU状态和信息
nvidia-smi
定时刷新状态,-n(秒),每秒刷新一次
watch -n 1 nvidia-smi
指定可见GPU
终端跑程序指定GPU
CUDA_VISIBLE_DEVICES=0 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块同时跑
其他
训练后保存的模型参数,再次导入模型时,出现报错:
Missing key(s) in state_dict: ...
解决办法:
在导入参数前加上import torch.nn as nn
from torch.backends import cudnn
model = torch.nn.DataParallel(model)
cudnn.benchmark = True