1、模型的保存
在pytorch官网教程上很详细的举出示例,说明pytorch如何保存模型。
链接: pytorch保存模型示例
关于模型的保存可以分为以下两种,第一种是训练好模型进行保存,第二种是以字典的方式保存模型
(1)保存模型
torch.save(model.state_dict(), './model/model_weights.pth') #保存模型的所有参数
torch.save(model, './model/model_weights.pth') #保存模型
model.state_dict()模型调用这个函数,可以拿到模型所有的参数,而后面的参数则是保存的路径与模型。第一种不包含模型结构,所以模型占用内存小,第二种完整的保存了模型
(2)加载模型
model.load_state_dict(torch.load('./model/model_weights.pth')) #对模型中参数加载
torch.load('./model/model_weights.pth') #对模型进行加载
model.eval()
torch.load()函数则是torch.save()函数保存的模型,然后load_state_dict()加载进入模型的参数
model.eval()函数是转化为评估模式。不记录其中的一些值
2、字典的方式保存模型
(1)保存模型
注意:在保存模型的时候一定要把对于模型训练重要的东西添加进去,例如优化器,epoch,损失函数loss等。
# Additional information
EPOCH = 5
PATH = "./model/model_weights.pth"
LOSS = 0.4
torch.save({
'epoch': EPOCH,
'model_state_dict': net.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': LOSS,
}, PATH)
(2)加载模型
model = Net()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
checkpoint = torch.load(PATH) #加载进checkpoint这个字典中
model.load_state_dict(checkpoint['model_state_dict']) #对应字典传对应参数
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch'] #拿出来,在此基础上继续训练
loss = checkpoint['loss']
model.eval() # 或者 model.train()
3、保存多个模型
(1)保存模型
# Specify a path to save to
PATH = "./model/model_weights.pth"
torch.save({
'modelA_state_dict': netA.state_dict(),
'modelB_state_dict': netB.state_dict(),
'optimizerA_state_dict': optimizerA.state_dict(),
'optimizerB_state_dict': optimizerB.state_dict(),
}, PATH)
(2)加载模型
modelA = Net()
modelB = Net()
optimModelA = optim.SGD(modelA.parameters(), lr=0.001, momentum=0.9)
optimModelB = optim.SGD(modelB.parameters(), lr=0.001, momentum=0.9)
checkpoint = torch.load(PATH)
modelA.load_state_dict(checkpoint['modelA_state_dict'])
modelB.load_state_dict(checkpoint['modelB_state_dict'])
optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])
modelA.eval() #modelA.train()
modelB.eval() #modelB.train()