首页 > 解决方案 > 如何使用模型的预训练权重在下一次迭代中初始化权重?

问题描述

我有一个模型架构。torch.save()我已经使用n 次迭代保存了整个模型。我想通过使用我之前保存的模型的预训练权重来运行我的代码的另一次迭代。

编辑:我希望从预训练模型的权重中完成新迭代的权重初始化

编辑2:只是补充一点,我不打算恢复训练。我打算保存模型并将其用于具有相同参数的单独训练。可以将其想象为使用带有权重等的已保存模型来进行更大的运行和更多的样本(即全新的训练工作)

现在,我做类似的事情:

# default_lr = 5
# default_weight_decay = 0.001
# model_io = the pretrained model 
model = torch.load(model_io) 
optim = torch.optim.Adam(model.parameters(),lr=default_lr, weight_decay=default_weight_decay)  
loss_new = BCELoss()  
epochs = default_epoch 
.
.
training_loop():
....
outputs = model(input)
....
.
#similarly for test loop

我错过了什么吗?我必须为大量样本运行很长时间,所以不能等着看结果然后弄清楚事情。

谢谢!

标签: machine-learningdeep-learningpytorchtorchvision

解决方案


从您发布的代码中,我看到您只是加载以前的模型参数,以便从您离开的地方重新开始训练。这不足以正确地重新开始训练。除了模型参数(权重)之外,您还需要保存和加载优化器状态,​​尤其是当您选择的优化器是 Adam 时,它具有所有权重的速度参数,有助于降低学习率。

为了顺利重新开始训练,我会做以下事情:

# For saving your model

state = {
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict()
}
model_save_path = "Enter/your/model/path/here/model_name.pth"
torch.save(state, model_save_path)

# ------------------------------------------

# For loading your model
state = torch.load(model_save_path)

model = MyNetwork()
model.load_state_dict(state['model'])

optim = torch.optim.Adam(model.parameters(),lr=default_lr, weight_decay=default_weight_decay)
optim.load_state_dict(state['optimizer'])

除此之外,如果您使用学习率衰减策略,您可能还希望保存学习率、迄今为止您可能需要用于检查点目的的最佳验证准确度,以及可能影响训练的任何其他可变参数。但在大多数情况下,仅保存和加载模型权重和优化器状态就足够了。

编辑:您可能还想查看以下答案,该答案详细解释了如何在不同情况下保存模型。


推荐阅读