PyTorch模型保存的两种方式

在PyTorch中，保存模型主要有两种方式：使用torch.save()和torch.jit.script()。这两种方式各有优缺点，适用于不同的场景。
1. 使用torch.save()保存模型
torch.save()是PyTorch中用于保存模型的标准方法。你可以将模型的所有参数和缓冲区保存到硬盘上，以便以后使用torch.load()重新加载。
示例代码：

import torch
# 假设你已经训练好了一个模型，命名为model
model = ...
# 使用torch.save()保存模型
torch.save(model.state_dict(), 'model_weights.pth')

在上面的代码中，model.state_dict()返回一个包含模型参数的字典，该字典随后被保存到名为model_weights.pth的文件中。当你需要重新加载模型时，可以使用torch.load()函数。
2. 使用torch.jit.script()保存模型
torch.jit.script()是另一种保存模型的方法，它可以将整个模型（包括模型结构）保存为一个TorchScript脚本。TorchScript脚本可以在不依赖Python解释器的情况下运行，这使得模型部署更加方便。
示例代码：

import torch
# 假设你已经训练好了一个模型，命名为model
model = ...
# 使用torch.jit.script()保存模型为TorchScript脚本
traced_script_module = torch.jit.script(model)
traced_script_module.save('model_script.pt')

在上面的代码中，torch.jit.script(model)将模型转换为TorchScript脚本，然后使用traced_script_module.save('model_script.pt')将其保存到硬盘上。之后，你可以使用torch.jit.load()重新加载模型。
这两种方法各有优缺点。使用torch.save()保存模型只包含模型的参数，而使用torch.jit.script()保存模型则包括整个模型结构。这意味着使用torch.jit.script()可以更方便地部署模型，因为它不需要Python解释器。然而，torch.save()可以让你更容易地更改和修改模型结构，而无需重新训练或转换模型。因此，选择哪种方法取决于你的具体需求。如果你需要部署模型并使其在没有Python解释器的环境中运行，那么使用torch.jit.script()是一个更好的选择。如果你需要经常修改和调整模型结构，那么使用torch.save()可能更加方便。