pytorch如何保存训练模型参数并实现继续训练
最近的想法是在推荐模型中考虑根据用户对推荐结果的后续选择，利用已训练的offline预训练模型参数来更新新的结果。
简单记录一下中途保存参数和后续使用不同数据训练的方法。
简单模型和训练数据
先准备一个简单模型，简单两层linear出个分类结果。
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(64, 32)
self.linear1 = nn.Linear(32, 10)
self.relu = nn.ReLU()
def forward(self, x):
x = self.linear(x)
x = self.relu(x)
x = self.linear1(x)
return x准备训练用数据，这里直接随机两份，同时给出配套的十个类的分类label结果。
要注意的是 crossEntropy 交叉熵只认 long 以上的tensor，label记得转一下类型。
rand1 = torch.rand((100, 64)).to(torch.float)
label1 = np.random.randint(0, 10, size=100)
label1 = torch.from_numpy(label1).to(torch.long)
rand2 = torch.rand((100, 64)).to(torch.float)
label2 = np.random.randint(0, 10, size=100)
label2 = torch.from_numpy(label2).to(torch.long)训练简单使用交叉熵，优化器Adam。
model = MyModel()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss = nn.CrossEntropyLoss()
iteration = 100
for i in range(iteration):
output = model(rand1)
my_loss = loss(output, label1)
optimizer.zero_grad()
my_loss.backward()
optimizer.step()
print("iteration:{} loss:{}".format(i, my_loss))反正能跑起来：
保存与读取训练参数结果的方法
关键的保存方法，可以分为两种，一种是直接把模型整体保存：
torch.save(model, save_path)
两个参数，模型和保存目录。不过这种不常用，如果模型变化或者只需要其中一部分参数就不太灵活。
常用方法的是将需要的模型或优化器参数取出以字典形式存储，这样可以在使用时初始化相关模型，读入对应参数即可。
def save_model(save_path, iteration, optimizer, model):
torch.save({'iteration': iteration,
'optimizer_dict': optimizer.state_dict(),
'model_dict': model.state_dict()},
save_path)
print("model save success")分别存储训练循环次数，优化器设置和模型参数结果。
初始化模型，读取参数并设置：
def load_model(save_name, optimizer, model):
model_data = torch.load(save_name)
model.load_state_dict(model_data['model_dict'])
optimizer.load_state_dict(model_data['optimizer_dict'])
print("model load success")初始化新模型：
path = "net.dict"
save_model(path, iteration, optimizer, model)
print(model.state_dict()['linear.weight'])
new_model = MyModel()
new_optimizer = torch.optim.Adam(new_model.parameters(), lr=0.01)
load_model(path, new_optimizer, new_model)
print(new_model.state_dict()['linear.weight'])输出第一个linear层的参数看看，确实相同，参数成功读取上了。注意optimizer的初始化对应模型别写错了。
之后用新模型继续训练试试：
for i in range(iteration):
output = new_model(rand2)
my_loss = loss(output, label2)
new_optimizer.zero_grad()
my_loss.backward()
new_optimizer.step()
print("iteration:{} loss:{}".format(i, my_loss))能成功训练。
变化学习率的保存
上面的demo只用了固定的学习率来做实验。
如果使用了 scheduler 来变化步长，只要保存 scheduler 的 state_dict ，之后对新初始化的 scheduler 设置对应的当前循环次数即可。
# 存储时 'scheduler': scheduler.state_dict() # 读取时 scheduler.load_state_dict(checkpoint['lr_schedule'])
scheduler的使用可以看看我之前整理的文章：利用scheduler实现learning-rate学习率动态变化
总结
这次主要是整理了一下pytorch模型参数的整体保存方法，来实现新数据的后续训练或直接作为offline预训练模型来使用。
不过后续数据分布不同的话感觉效果会很差啊…
也不知道能不能用什么算法修改下权重来贴合新的数据，找点多次训练优化论文看看好了。
以上为个人经验，希望能给大家一个参考，也希望大家多多支持脚本之家。
相关文章
这篇文章主要介绍了Python读写文件模式和文件对象方法,结合实例形式详细分析了Python文件操作常用技巧与相关注意事项,需要的朋友可以参考下2019-09-09
这篇文章主要介绍了python 如何将IP 地址转整数，帮助大家了解转换的原理与收益，更好的理解python，感兴趣的朋友可以了解下2020-11-11
这篇文章主要为大家详细介绍了如何使用Python Tkinter编写一个学生姓名比较工具,文中的示例代码讲解详细,感兴趣的小伙伴可以跟随小编一起学习一下2025-10-10
python中np.random.permutation函数实例详解
np.random.permutation是numpy中的一个函数,它可以将一个数组中的元素随机打乱,返回一个打乱后的新数组,下面这篇文章主要给大家介绍了关于python中np.random.permutation函数的相关资料,需要的朋友可以参考下2023-04-04
这篇文章主要为大家详细介绍了如何利用Python语言制作简易的核酸检测日历，文中的示例代码讲解详细，感兴趣的小伙伴可以动手尝试一下2022-09-09
今天小编就为大家分享一篇Python自定义一个类实现字典dict功能的方法，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧2019-01-01
这篇文章主要介绍了Django中使用Celery的方法示例，小编觉得挺不错的，现在分享给大家，也给大家做个参考。一起跟随小编过来看看吧2018-11-11
Python常见读写文件操作实例总结【文本、json、csv、pdf等】
这篇文章主要介绍了Python常见读写文件操作,结合实例形式总结分析了Python常见的各种文件读写操作,包括文本、json、csv、pdf等文件的读写与相关注意事项,需要的朋友可以参考下2019-04-04
这篇文章主要介绍了详解Django的CSRF认证实现，详细的介绍了csrf原理和实现，小编觉得挺不错的，现在分享给大家，也给大家做个参考。一起跟随小编过来看看吧2018-10-10
在日常数据处理的过程中，有时候拿到手的是完整的地址信息，如果需要从地址信息中解析出相应的各级行政单元名称，方式有很多，而今天要介绍的方式只需要一行代码即可快速实现，快跟随小编一起学习一下吧2022-05-05
最新评论