4. 模型序列化与反序列化

任务说明

训练好的PyTorch模型需要部署到生产环境时，通常要脱离训练框架或跨语言调用。本案例将完成：把训练好的模型分别导出为ONNX和TorchScript格式，再加载这两种格式执行推理，验证输出与原始PyTorch模型一致，从而确认跨框架兼容性。

环境依赖

• Python 3.8+
• torch 2.0+
• onnxruntime 1.15+（用于加载ONNX并推理）

完整实现

import torch
import torch.nn as nn
import onnxruntime as ort
import numpy as np

# 1. 定义并训练一个简单模型（用固定数据模拟训练完成）
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(4, 2)

    def forward(self, x):
        return torch.softmax(self.fc(x), dim=1)

model = SimpleNet()
model.eval()

# 模拟训练后的权重（实际场景中此处为训练好的模型）
with torch.no_grad():
    model.fc.weight.copy_(torch.tensor([[0.5, -0.2, 0.1, 0.3],
                                        [-0.4, 0.6, 0.2, -0.1]]))
    model.fc.bias.copy_(torch.tensor([0.1, -0.2]))

# 2. 导出为ONNX格式
dummy_input = torch.randn(1, 4)
torch.onnx.export(model, dummy_input, "simple_model.onnx",
                  input_names=["input"], output_names=["output"],
                  dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}})

# 3. 导出为TorchScript格式
scripted_model = torch.jit.script(model)
scripted_model.save("simple_model.pt")

# 4. 加载ONNX并推理
ort_session = ort.InferenceSession("simple_model.onnx")
test_input = np.array([[0.8, -0.5, 0.2, 1.0]], dtype=np.float32)
onnx_output = ort_session.run(None, {"input": test_input})[0]

# 5. 加载TorchScript并推理
loaded_script = torch.jit.load("simple_model.pt")
ts_output = loaded_script(torch.from_numpy(test_input)).detach().numpy()

# 6. 与原始PyTorch模型输出对比
with torch.no_grad():
    original_output = model(torch.from_numpy(test_input)).numpy()

print("ONNX输出:", onnx_output)
print("TorchScript输出:", ts_output)
print("原始模型输出:", original_output)

# 验证一致性
assert np.allclose(onnx_output, original_output, atol=1e-6), "ONNX输出不一致"
assert np.allclose(ts_output, original_output, atol=1e-6), "TorchScript输出不一致"
print("验证通过：两种格式输出与原始模型一致")

运行方式

pip install torch onnxruntime
python serialize_model.py

验证方法与判定标准

运行脚本后，控制台会打印三种输出。判定通过的条件：

• 程序无异常退出，且打印“验证通过”信息；
• ONNX和TorchScript的输出与原始PyTorch模型输出数值一致（误差小于1e-6）；
• 当前目录下生成simple_model.onnx和simple_model.pt两个文件，说明导出成功。

关键说明

ONNX导出时通过dynamic_axes指定了动态batch维度，使模型可接受不同batch大小的输入。TorchScript使用torch.jit.script进行静态编译，适合无动态控制流的模型。若模型包含数据相关分支，应改用torch.jit.trace。ONNX文件可通过onnxruntime在C++、Java等环境中加载，TorchScript则依赖LibTorch，两者均实现了跨框架部署能力。