一、假设我们只保存了模型的参数(model.state_dict())到文件名为modelparameters.pth, model = Net()
1. cpu -> cpu或者gpu -> gpu:
checkpoint = torch.load('modelparameters.pth')
model.load_state_dict(checkpoint)
2. cpu -> gpu 0
checkpoint = torch.load('modelparameters.pth', map_location=lambda storage, loc: storage.cuda(0))
model.load_state_dict(checkpoint)
3. gpu 0 -> gpu 1
checkpoint = torch.load('modelparameters.pth', map_location={'cuda:1':'cuda:0'})
model.load_state_dict(checkpoint)
4. gpu -> cpu
checkpoint = torch.load('modelparameters.pth', map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint)
二、pytorch完整加载模型实例
# 模型是GPU训练的
def load_models(model_param_path):
torch.manual_seed(14)
np.random.seed(14)
random.seed(14)
if torch.cuda.is_available():
cudnn.benchmark = True
torch.cuda.manual_seed_all(14)
model = FeatherNetB()
model = torch.nn.DataParallel(model)
device = torch.device('cuda:' + str(0) if torch.cuda.is_available() else "cpu")
model.to(device)
print("=> loading checkpoint '{}'".format(param_path))
if torch.cuda.is_available():
#gpu测试
checkpoint = torch.load(model_param_path)
else:
#cpu测试
checkpoint = torch.load(model_param_path, map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint['state_dict'])
#测试的时候必须加上model.eval(),训练就用model.train()
model.eval()
return model