我正在寻找在PyTorch中保存训练模型的替代方法。到目前为止，我找到了两种选择。
我已经遇到过这个discussion，其中推荐方法2超过方法1。
我的问题是，为什么第二种方法更受欢迎？是否因为torch.nn模块具有这两个功能，我们被鼓励使用它们？
我在他们的github repo上找到了this page，我只是在这里粘贴内容。
序列化和恢复模型有两种主要方法。
第一个（推荐）保存并仅加载模型参数：
torch.save(the_model.state_dict(), PATH)
然后呢：
the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))
第二个保存并加载整个模型：
torch.save(the_model, PATH)
然后呢：
the_model = torch.load(PATH)
但是在这种情况下，序列化数据绑定到特定的类和使用的确切目录结构，因此当在其他项目中使用时，或者在一些严重的重构之后，它可以以各种方式中断。
这取决于你想做什么。
案例＃1：保存模型以自行使用它进行推理：保存模型，恢复模型，然后将模型更改为评估模式。这样做是因为你通常有BatchNorm和Dropout图层，默认情况下在构造中处于训练模式：
torch.save(model.state_dict(), filepath)
#Later to restore:
model.load_state_dict(torch.load(filepath))
model.eval()
案例＃2：保存模型以便以后恢复训练：如果您需要继续训练您将要保存的模型，则需要保存的不仅仅是模型。您还需要保存优化器，时期，分数等的状态。您可以这样做：
state = {
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
...
}
torch.save(state, filepath)
要恢复训练，您可以执行以下操作：state = torch.load(filepath)，然后，恢复每个对象的状态，如下所示：
model.load_state_dict(state['state_dict'])
optimizer.load_state_dict(state['optimizer'])
由于您正在恢复训练，因此在加载时恢复状态时请勿调用model.eval()。
案例＃3：其他人无法访问您的代码使用的模型：在Tensorflow中，您可以创建一个.pb文件，该文件定义了模型的体系结构和权重。这非常方便，特别是在使用Tensorflow serve时。在Pytorch中执行此操作的等效方法是：
torch.save(model, filepath)
# Then later:
model = torch.load(filepath)
这种方式仍然不是防弹，因为pytorch仍然经历了很多变化，我不推荐它。
pickle Python库实现了用于序列化和反序列化Python对象的二进制协议。
当你import torch（或当你使用PyTorch）它将为你import pickle，你不需要直接调用pickle.dump()和pickle.load()，这是保存和加载对象的方法。
事实上，torch.save()和torch.load()将为你包装pickle.dump()和pickle.load()。
一个state_dict提到的另一个答案值得更多的笔记。
我们在PyTorch中有什么state_dict？实际上有两个state_dicts。
PyTorch模型是torch.nn.Module有model.parameters()调用以获得可学习的参数（w和b）。这些可学习的参数，一旦随机设置，将在我们学习时随时间更新。可学习的参数是第一个state_dict。
第二个state_dict是优化器状态字典。您记得优化器用于改善我们的可学习参数。但优化器state_dict是固定的。那里没什么好学的。
因为state_dict对象是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)
print("Model's state_dict:")
for param_tensor in model.state_dict():
print(param_tensor, "\t", model.state_dict()[param_tensor].size())
print("Model weight:")
print(model.weight)
print("Model bias:")
print(model.bias)
print("---")
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
print(var_name, "\t", optimizer.state_dict()[var_name])
此代码将输出以下内容：
Model's state_dict:
weight torch.Size([2, 5])
bias torch.Size([2])
Model weight:
Parameter containing:
tensor([[ 0.1328, 0.1360, 0.1553, -0.1838, -0.0316],
[ 0.0479, 0.1760, 0.1712, 0.2244, 0.1408]], requires_grad=True)
Model bias:
Parameter containing:
tensor([ 0.4112, -0.0733], requires_grad=True)
---
Optimizer's state_dict:
state {}
param_groups [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [140695321443856, 140695321443928]}]
请注意，这是一个最小的模型。您可以尝试添加顺序堆栈
model = torch.nn.Sequential(
torch.nn.Linear(D_in, H),
torch.nn.Conv2d(A, B, C)
torch.nn.Linear(H, D_out),
)
请注意，只有具有可学习参数（卷积层，线性层等）和注册缓冲区（batchnorm图层）的图层在模型的state_dict中具有条目。
不可学习的东西，属于优化器对象state_dict，它包含有关优化器状态的信息，以及使用的超参数。
其余的故事是一样的;在推理阶段（这是我们在训练后使用模型的阶段）进行预测;我们根据我们学到的参数进行预测。所以对于推理，我们只需要保存参数model.state_dict()。
torch.save(model.state_dict(), filepath)
并使用稍后的model.load_state_dict（torch.load（filepath））model.eval（）
注意：不要忘记最后一行model.eval()这在加载模型后至关重要。
也不要试图保存torch.save(model.parameters(), filepath)。 model.parameters()只是生成器对象。
另一方面，torch.save(model, filepath)保存模型对象本身，但请记住，模型没有优化器的state_dict。查看@Jadiel de Armas的另一个优秀答案，以保存优化器的状态字典。
常见的PyTorch约定是使用.pt或.pth文件扩展名保存模型。
保存/加载整个模型保存：
path = "username/directory/lstmmodelgpu.pth"
torch.save(trainer, path)
加载：
model = torch.load(PATH)
model.eval()