3. 模型导出与序列化：PyTorch 模型保存与加载、ONNX 导出、TorchScript 转换

GraphSAGE 模型从训练环境迁移到生产推理系统，核心挑战在于 序列化格式的兼容性 与 推理性能的优化。本章聚焦三种主流方案：PyTorch 原生序列化、ONNX 跨平台导出、TorchScript 静态图编译。每种方案在部署场景、算子支持度、推理速度上存在显著差异，需根据上线环境做出选择。

3.1 PyTorch 模型保存与加载

PyTorch 提供两种序列化范式：状态字典（state_dict） 与 完整模型（torch.save）。对于 GraphSAGE 这类包含自定义 SAGEConv 层的模型，推荐使用 state_dict 方式，以保持代码与权重的解耦。

# 保存
torch.save({
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'epoch': epoch,
    'args': {'in_dim': 1433, 'hidden_dim': 128, 'out_dim': 7, 'aggregator': 'mean'}
}, "gsage_checkpoint.pt")

# 加载
checkpoint = torch.load("gsage_checkpoint.pt", map_location='cpu')
model = GraphSAGE(checkpoint['args']['in_dim'],
                  checkpoint['args']['hidden_dim'],
                  checkpoint['args']['out_dim'],
                  aggregator=checkpoint['args']['aggregator'])
model.load_state_dict(checkpoint['model_state_dict'])

3.2 ONNX 导出：跨平台与硬件加速

ONNX（Open Neural Network Exchange）为 GraphSAGE 提供 框架无关的中间表示，可部署到 ONNX Runtime、TensorRT、OpenVINO 等推理引擎。导出时需处理图神经网络特有的 可变输入形状 与 稀疏邻接矩阵。

import torch.onnx

# 准备固定形状的示例输入
dummy_x = torch.randn(1, 1433)          # 单节点特征
dummy_adj = torch.randn(1, 1433, 1433)  # 稠密邻接矩阵（仅用于导出）

torch.onnx.export(
    model,
    (dummy_x, dummy_adj),
    "graphsage.onnx",
    input_names=['node_feat', 'adj_matrix'],
    output_names=['embedding'],
    dynamic_axes={
        'node_feat': {0: 'batch_size'},
        'adj_matrix': {0: 'batch_size'},
        'embedding': {0: 'batch_size'}
    },
    opset_version=17
)

3.3 TorchScript 转换：静态图与 C++ 部署

TorchScript 将动态 PyTorch 模型编译为 静态计算图，消除 Python 解释器开销，支持 C++ 生产环境直接加载。GraphSAGE 的转换有两种路径：追踪（Tracing） 与 脚本化（Scripting）。

• Tracing（torch.jit.trace）：通过一次前向传播记录算子序列，适合不含数据依赖控制流的模型。GraphSAGE 的固定深度采样场景适用。
• Scripting（torch.jit.script）：直接编译 Python 代码，支持 if/for 等控制流，但要求模型代码为 TorchScript 子集（无动态字典、无外部库调用）。

# 方法一：Tracing（推荐用于推理图固定的场景）
traced_model = torch.jit.trace(model, (dummy_x, dummy_adj))
traced_model.save("graphsage_traced.pt")

# 方法二：Scripting（需修改模型代码适配 TorchScript）
@torch.jit.script
def sage_forward(x, adj, weight):
    # 自定义聚合逻辑
    return torch.mm(adj, torch.mm(x, weight))

scripted_model = torch.jit.script(model)
scripted_model.save("graphsage_scripted.pt")

3.4 三种方案对比与选型

3.5 序列化验证清单

模型导出后必须通过 数值一致性验证，确保推理结果与原始 PyTorch 模型在相同输入下的误差在可接受范围内（通常 < 1e-5）。

def validate_export(original_model, exported_path, dummy_input):
    # 原始模型推理
    original_model.eval()
    with torch.no_grad():
        ref_output = original_model(*dummy_input)

    # 加载导出模型
    if exported_path.endswith('.onnx'):
        import onnxruntime as ort
        session = ort.InferenceSession(exported_path)
        ort_inputs = {session.get_inputs()[i].name: dummy_input[i].numpy()
                      for i in range(len(dummy_input))}
        exported_output = torch.tensor(session.run(None, ort_inputs)[0])
    else:  # TorchScript
        exported_model = torch.jit.load(exported_path)
        exported_output = exported_model(*dummy_input)

    # 计算最大相对误差
    error = torch.max(torch.abs(ref_output - exported_output) /
                      (torch.abs(ref_output) + 1e-8))
    assert error < 1e-5, f"导出验证失败，最大相对误差：{error:.2e}"
    print(f"✅ 导出验证通过，最大相对误差：{error:.2e}")