torch.save

torch.save函数是PyTorch中用于保存模型的函数。一般约定使用.pt或.pth文件扩展名保存模型。该函数的实现在torch/serialization.py文件中。[1] 保存模型时，除了保存模型的state_dict外，还可以保存优化器的state_dict以及其他相关信息，如已训练的epoch编号、最新记录的训练损失等。这样的保存通常比单独保存模型要大2至3倍，因为它包含了额外的信息。[2] 当使用torch.load函数加载模型时，需要注意load_state_dict函数需要传递一个字典对象，而不是保存对象的路径。因此，在调用load_state_dict函数之前，需要对保存的state_dict进行反序列化操作。[3]

相关问题

torch.save和torch.jit.save

`torch.save` 和 `torch.jit.save` 都是PyTorch提供的用于保存模型和其相关状态的方法，但在用途上有所不同：

1. **torch.save**：
这是一个通用的模型保存函数，用于保存整个神经网络模型及其状态（例如权重、优化器状态等）。当你使用 `torch.save(obj, file_path)` 时，`obj` 可以是`nn.Module`实例（模型），`state_dict` 或者包含模型和数据的状态字典。它将模型及状态信息存储为二进制文件，便于后续加载恢复到相同的配置。如果你只想保存模型结构而忽略运行时计算图，可以只保存`state_dict`。

   state_dict = model.state_dict()  # 模型的权重
   torch.save(state_dict, 'model.pth')  # 保存权重

2. **torch.jit.save**：
`torch.jit.save` 是用于保存PyTorch的静态图（即通过`torch.jit.trace`或`torch.jit.script`生成的）模型，这种模型可以在没有原始Python代码的情况下独立运行。它可以保存模型的计算图，包括输入类型的元数据，使得其他程序可以像使用常规模块一样调用它，无需完整的Python环境。

   traced_model = torch.jit.trace(model, input_image)  # 使用输入图像trace模型
   torch.jit.save(traced_model, 'traced_model.pt')  # 保存为静态图模型

使用`torch.jit.save`保存的模型通常比`torch.save`更小，因为它不包含Python代码，而且加载速度也更快，因为它是静态执行的。

torch.export.save和torch.save

`torch.save` 是 PyTorch 中用于序列化和持久化模型及张量的函数。它可以将一个 Python 对象保存到硬盘上，对象通常是一个 PyTorch `Tensor`，或者是一个模型对象（即一个包含可训练参数的 `nn.Module` 类实例）。保存的对象可以使用 `torch.load` 进行反序列化，这样就可以在之后重新加载模型或张量到内存中。

import torch

# 保存张量
tensor = torch.tensor([1, 2, 3])
torch.save(tensor, 'tensor.pt')

# 保存模型
model = torch.nn.Linear(3, 4)
torch.save(model.state_dict(), 'model_weights.pt')

`torch.export.save` 不是 PyTorch 的一个内置函数。可能你指的是 `torch.save` 或者是 PyTorch 的导出功能（例如 TorchScript 或者 ONNX），这些功能用于将模型转换为可以在不同环境中运行的格式。例如，TorchScript 允许将模型转换为 TorchScript 格式，这样就可以在没有 Python 依赖的环境中运行模型。

# 使用 TorchScript 导出模型
model = torch.jit.trace(model, example_input)
model.save('model_scripted.pt')

或者，使用 ONNX 导出模型，使其可以在支持 ONNX 的推理引擎上运行：

# 导出模型为 ONNX 格式
input_sample = torch.randn((1, 3, 224, 224))
torch.onnx.export(model, input_sample, "model.onnx")

在使用这些功能时，重要的是要理解你正在导出的模型需要在什么环境下运行，以及模型的输入输出接口是否与导出格式兼容。