用火把拯救一个受过训练的模特的最佳方法？

Question

我在寻找另一种方法来拯救一个在PyTorch受过训练的模特。到目前为止，我已经找到了两种选择。

1. torch.save()保存模型，torch.load()加载模型。
2. model.state_dict()保存已训练的模型，model.load_state_dict()加载已保存的模型。

我遇到了这个discussion，这里推荐方法2而不是方法1。

我的问题是，为什么选择第二种方法？这仅仅是因为torch.nn模块有这两个功能，我们被鼓励使用它们吗？

Answer 1

一个常见的PyTorch约定是使用.pt或.pth文件扩展名保存模型。

保存/加载整个模型 保存：

path = "username/directory/lstmmodelgpu.pth"
torch.save(trainer, path)

加载：

必须在某处定义模型类

model = torch.load(PATH)
model.eval()