2. 模型导出与序列化：使用PyTorch导出GraphSAGE模型、TorchScript与ONNX格式的选择与转换
将训练完成的GraphSAGE模型从研究环境迁移到生产推理系统，核心步骤是模型导出与序列化。本章聚焦于PyTorch生态下两种主流序列化格式——TorchScript与ONNX，分析其原理、转换流程及在GraphSAGE场景下的选型策略。
核心目标：将动态计算图的PyTorch模型转换为静态、可移植、无Python依赖的序列化表示，为后续低延迟推理引擎（LibTorch、ONNX Runtime、TensorRT）提供输入。
2.1 导出前的模型准备：冻结与追踪
GraphSAGE模型包含可学习的聚合器（MeanPool、LSTM、GCN）和线性变换层。导出前需将模型置于eval()模式，并确保输入张量形状固定（或使用动态轴）。PyTorch提供两种导出路径：
- 追踪（Tracing）： 通过一次前向传播记录计算图，适用于不含数据依赖控制流的模型。GraphSAGE的聚合逻辑通常为固定结构，适合追踪。
- 脚本化（Scripting）： 编译完整的Python代码为TorchScript IR，适用于包含
if/for等动态分支的模型。
import torch
import torch.nn as nn
class GraphSAGE(nn.Module):
def __init__(self, in_feats, hidden_feats, out_feats):
super().__init__()
self.conv1 = SAGEConv(in_feats, hidden_feats, aggregator_type='mean')
self.conv2 = SAGEConv(hidden_feats, out_feats, aggregator_type='mean')
def forward(self, x, edge_index):
x = self.conv1(x, edge_index).relu()
x = self.conv2(x, edge_index)
return x
model = GraphSAGE(128, 64, 32).eval()
dummy_x = torch.randn(1000, 128)
dummy_edge = torch.randint(0, 1000, (2, 5000))
# 方法一：TorchScript 追踪
traced_model = torch.jit.trace(model, (dummy_x, dummy_edge))
traced_model.save('graphsage_traced.pt')
# 方法二：TorchScript 脚本化（如需动态控制流）
scripted_model = torch.jit.script(model)
scripted_model.save('graphsage_scripted.pt')
实践建议：对于标准GraphSAGE（MeanPool聚合），优先使用
torch.jit.trace。若模型包含torch.where或基于节点度的分支逻辑，则必须使用torch.jit.script。
2.2 TorchScript：PyTorch原生序列化格式
TorchScript是PyTorch的原生中间表示，可直接被LibTorch（C++推理库）加载，无需Python解释器。其优势在于与PyTorch算子集完全对齐，支持自动微分和梯度计算（推理时通常关闭）。
| 特性 | TorchScript | 适用场景 |
|---|---|---|
| 算子覆盖度 | 100% PyTorch算子 | 自定义算子、复杂控制流 |
| 推理引擎 | LibTorch (C++) | 需要PyTorch生态深度集成 |
| 图优化 | 中等（融合、常量折叠） | 对延迟敏感度中等 |
| 跨框架兼容 | 仅PyTorch | 单一技术栈团队 |
2.3 ONNX：跨框架互操作标准
ONNX（Open Neural Network Exchange）提供框架无关的模型表示。将GraphSAGE导出为ONNX后，可被ONNX Runtime、TensorRT、OpenVINO等推理后端加速。但需注意：图神经网络中的稀疏操作（如邻域采样、消息传递）可能不被ONNX原生支持。
import torch.onnx
# 导出ONNX（需指定动态轴以支持可变批量）
torch.onnx.export(
model, # 模型
(dummy_x, dummy_edge), # 示例输入
'graphsage.onnx', # 输出路径
input_names=['x', 'edge_index'],
output_names=['output'],
dynamic_axes={
'x': {0: 'batch_size'},
'edge_index': {1: 'num_edges'},
'output': {0: 'batch_size'}
},
opset_version=17, # 推荐 ≥16 以支持更多算子
do_constant_folding=True
)
常见陷阱：GraphSAGE中的
scatter_add或segment_sum等聚合操作在ONNX中可能映射为ReduceSum或自定义算子。导出前需检查算子兼容性，必要时使用torch.onnx.register_custom_op_symbolic注册自定义映射。
2.4 TorchScript vs ONNX：选型对比
| 对比维度 | TorchScript | ONNX |
|---|---|---|
| 导出复杂度 | 低（一行代码） | 中（需处理算子映射） |
| 图神经网络支持 | 原生支持所有PyTorch GNN算子 | 部分稀疏算子需自定义 |
| 推理延迟（GPU） | 中等（LibTorch优化有限） | 低（TensorRT/ONNX Runtime深度优化） |
| 跨平台部署 | 仅PyTorch/C++生态 | 多框架、多硬件（CPU/GPU/NPU） |
| 动态形状支持 | 良好（需设置dynamic_axes） | 良好（ONNX Runtime支持） |
2.5 转换验证与精度对齐
导出后必须验证序列化模型的输出与原始PyTorch模型一致。推荐使用余弦相似度或最大绝对误差作为指标。
import numpy as np
# 加载导出的模型
traced = torch.jit.load('graphsage_traced.pt')
onnx_model = onnx.load('graphsage.onnx')
onnx.checker.check_model(onnx_model)
# 推理对比
with torch.no_grad():
ref_out = model(dummy_x, dummy_edge)
traced_out = traced(dummy_x, dummy_edge)
# 计算误差
cos_sim = torch.nn.functional.cosine_similarity(
ref_out.flatten(), traced_out.flatten(), dim=0
)
max_err = torch.max(torch.abs(ref_out - traced_out))
print(f'Cosine Similarity: {cos_sim.item():.6f}')
print(f'Max Absolute Error: {max_err.item():.6f}')
# 对于ONNX，使用ONNX Runtime验证
import onnxruntime as ort
ort_session = ort.InferenceSession('graphsage.onnx')
onnx_out = ort_session.run(
['output'],
{'x': dummy_x.numpy(), 'edge_index': dummy_edge.numpy()}
)
onnx_out = torch.tensor(onnx_out[0])
cos_sim_onnx = torch.nn.functional.cosine_similarity(
ref_out.flatten(), onnx_out.flatten(), dim=0
)
print(f'ONNX Cosine Similarity: {cos_sim_onnx.item():.6f}')
精度阈值：对于FP32模型，余弦相似度应 > 0.9999，最大绝对误差 < 1e-5。若误差过大，检查是否误用了
model.train()模式或Dropout/BatchNorm层未冻结。
2.6 序列化最佳实践清单
- 固定输入形状：导出时使用代表性输入（节点数、边数接近生产环境），避免极端动态形状导致图优化失效。
- 算子版本对齐：ONNX opset_version 推荐 ≥ 16，TorchScript 使用与PyTorch版本匹配的IR版本。
- 常量折叠：导出时开启
do_constant_folding=True，将固定参数（如权重）预计算为常量。 - 批量推理适配：若生产环境使用批量推理，在导出时设置
dynamic_axes为批量维度。 - 自定义算子注册：对于GraphSAGE中的自定义聚合器，提前在ONNX中注册符号函数，避免导出失败。