python - PyTorch模型保存错误: "Can' t pickle local object"

Question

当我尝试用这段代码保存 PyTorch 模型时:

checkpoint = {'model': Net(), 'state_dict': model.state_dict(),'optimizer' :optimizer.state_dict()}
torch.save(checkpoint, 'Checkpoint.pth')

我收到以下错误:

    E:\PROGRAM FILES\Anaconda\envs\staj_projesi\lib\site-packages\torch\serialization.py:251: UserWarning: Couldn't retrieve source code for container of type Net. It won't be checked for correctness upon loading.
...

      "type " + obj.__name__ + ". It won't be checked "
    Can't pickle local object 'trainModel.<locals>.Net'

当我尝试用这段代码保存 PyTorch 模型时:

checkpoint = {'state_dict': model.state_dict(),'optimizer' :optimizer.state_dict()}
torch.save(checkpoint, 'Checkpoint.pth')

我没有收到任何错误，但我想保存 ANN 类。我怎么解决这个问题？另外，我可以在之前的其他项目中保存具有第一个结构的模型

Answer 1

你不能! torch.save 仅保存对象 state_dict()。

当您使用以下内容时:

checkpoint = {'model': Net(), 'state_dict': model.state_dict(),'optimizer' :optimizer.state_dict()}
torch.save(checkpoint, 'Checkpoint.pth')

您正在尝试保存模型本身，但此数据保存在 model.state_dict() 中，并且在使用 state_dict 加载模型时，您应该首先启动模型对象。

这正是第二种方法正常工作的原因:

checkpoint = {'state_dict': model.state_dict(),'optimizer' :optimizer.state_dict()}
torch.save(checkpoint, 'Checkpoint.pth')

我建议在以下链接中阅读有关如何正确保存\加载模型的 pytorch 文档: https://pytorch.org/tutorials/beginner/saving_loading_models.html