如何使用模型的预训练权重在下一次迭代中初始化权重?

时间:2020-06-02 01:18:35

标签: machine-learning deep-learning pytorch torchvision

我有一个模型架构。我已经使用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

我错过了什么吗?我必须运行很长时间才能获取大量样本,因此等不及要看结果然后再弄清楚。

谢谢!

1 个答案:

答案 0 :(得分:0)

从您发布的代码中,我看到您只是在加载先前的模型参数,以便从中断的地方重新开始训练。这不足以正确地重新开始训练。除了模型参数(权重)之外,还需要保存和加载优化器状态,​​尤其是当您选择的优化器是Adam时,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'])

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

编辑:您可能还想看看answer之后的内容,它详细说明了如何在不同情况下保存模型。