python如何保存训练好的模型_python - 在PyTorch中保存训练模型的最佳方法？

pickle Python库实现了用于序列化和反序列化Python对象的二进制协议。

当你state_dict（或当你使用PyTorch）它将为你torch.load()，你不需要直接调用pickle.dump()和pickle.load()，这是保存和加载对象的方法。

事实上，state_dict和torch.load()将为您包装pickle.dump()和pickle.load()。

一个state_dict提到的另一个答案值得再说几句。

我们在PyTorch中有什么torch.save(model, filepath)？实际上有两个state_dicts。

PyTorch模型是torch.save(model, filepath)具有state_dict调用以获得可学习的参数（w和b）。这些可学习的参数，一旦随机设置，将在我们学习时随时间更新。可学习的参数是第一个state_dict。

第二个torch.save(model, filepath)是优化器状态字典。 优化器也是该模型的一部分。 您记得优化器用于改善我们的可学习参数。 但是优化器state_dict是固定的。 那里没什么好学的。

因为torch.save(model, filepath)对象是Python字典，所以它们可以轻松保存，更新，更改和恢复，为PyTorch模型和优化器添加了大量模块化。

让我们创建一个超级简单的模型来解释这个：

import torch

import torch.optim as optim

model = torch.nn.Linear(5, 2)

# Initialize optimizer

optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9