一、指定一个GPU训练的两种方法:
1.代码中指定
import torch
torch.cuda.set_device(id)
2.终端中指定
CUDA_VISIBLE_DEVICES=1 python 你的程序
其中id就是你的gpu编号
二、多GPU并行训练:
torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)
该函数实现了在module级别上的数据并行使用,注意batch size要大于GPU的数量。
三、多GPU模型保存与加载
保存:
torch.save(model.module.state_dict(), model_out_path)
注意,要在模型名称的后面加上.module
加载:
from collections import OrderedDict
def myOwnLoad(model, check):
modelState = model.state_dict()
tempState = OrderedDict()
for i in range(len(check.keys())-2):
print modelState.keys()[i], check.keys()[i]
tempState[modelState.keys()[i]] = check[check.keys()[i]]
temp = [[0.02]*1024 for i in range(200)] # mean=0, std=0.02
tempState['myFc.weight'] = torch.normal(mean=0, std=torch.FloatTensor(temp)).cuda()
tempState['myFc.bias'] = torch.normal(mean=0, std=torch.FloatTensor([0]*200)).cuda()
model.load_state_dict(tempState)
return model