基于pytorch模型的保存与加载

最新推荐文章于 2024-07-25 14:35:42 发布

炸弹天堂

最新推荐文章于 2024-07-25 14:35:42 发布

阅读量226

点赞数 1

分类专栏： pytorch 文章标签： pytorch 深度学习人工智能

本文链接：https://blog.csdn.net/weixin_45932783/article/details/129301879

版权

pytorch 专栏收录该内容

4 篇文章 0 订阅

订阅专栏

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()