python - 在 pytorch 中训练期间 best_state 随模型变化

Question

我想保存最好的模型，然后在测试期间加载它。所以我使用了以下方法:

def train():  
    #training steps …  
    if acc > best_acc:  
        best_state = model.state_dict()  
        best_acc = acc
    return best_state  

然后，在主函数中我使用:

model.load_state_dict(best_state)  

恢复模型。

但是，我发现best_state始终与训练时的最后状态相同，而不是最佳状态。有谁知道原因以及如何避免它？

顺便说一句，我知道我可以使用torch.save(the_model.state_dict(), PATH)，然后通过以下方式加载模型the_model.load_state_dict(torch.load(PATH))。但是，我不想将参数保存到文件中，因为训练和测试函数位于一个文件中。

Answer 1

model.state_dict() 是 OrderedDict

from collections import OrderedDict

您可以使用:

from copy import deepcopy

解决问题

相反:

best_state = model.state_dict() 

您应该使用:

best_state = copy.deepcopy(model.state_dict())

深(而不是浅)复制使可变的 OrderedDict 实例不会改变 best_state 。

你可以查看我的other answer在 PyTorch 中保存状态字典。