PyTorch模型保存的两种方式：`torch.save()`与`torch.jit.script()`

在PyTorch中，模型保存主要有两种方式：使用torch.save()和torch.jit.script()。这两种方式各有优缺点，下面我们将分别介绍它们的特点和应用场景。

1. torch.save()

torch.save()是PyTorch中用于保存模型的最基本方法。你可以使用它将模型保存为.pt或.pth文件。使用这种方式保存的模型可以直接在PyTorch中加载和使用，无需进行额外的转换或编译。

import torch
# 假设我们有一个训练好的模型model
model = ...
# 使用torch.save()保存模型
torch.save(model.state_dict(), 'model.pt')

这种方式的主要优点是简单易用，适用于大多数情况。然而，它只保存了模型的参数，而没有保存模型的计算图。这意味着如果你更换了设备或环境，重新加载模型时可能需要重新编译计算图，这可能会影响加载速度和兼容性。

2. torch.jit.script()

torch.jit.script()可以将PyTorch模型转换为TorchScript格式，这是一种可以在不依赖Python解释器的环境中运行的中间表示形式。使用这种方式保存的模型具有更好的可移植性和兼容性，特别是对于那些需要在没有Python环境的设备上运行模型的情况。

import torch
# 假设我们有一个训练好的模型model
model = ...
# 使用torch.jit.script()将模型转换为TorchScript格式
traced_script_module = torch.jit.script(model)
# 保存TorchScript模型
traced_script_module.save('model.pt')

这种方式的主要优点是可移植性强，兼容性好。然而，它需要额外的转换步骤，而且只支持部分PyTorch功能。另外，由于TorchScript是一种中间表示形式，加载速度可能比直接使用torch.load()慢一些。

综上所述，选择哪种方式保存模型取决于你的具体需求。如果你需要将模型部署到没有Python环境的设备上，或者需要在不依赖原始训练环境的场景下运行模型，建议使用torch.jit.script()。否则，如果你只是需要在不同的设备或环境中重新加载模型，并且对加载速度要求不高，那么使用torch.save()可能更为简单方便。