4、模型转换：将PyTorch/TensorFlow模型转换为ONNX格式

4.1 核心原理：为什么需要“导出”而非“重写”

在Java部署ONNX模型之前，首先要理解模型转换的本质。PyTorch和TensorFlow是训练框架，它们内部维护的是动态计算图（PyTorch的Autograd）或静态图（TensorFlow的GraphDef），这些结构服务于训练过程中的梯度回传与参数更新。而ONNX（Open Neural Network Exchange）是一种推理交换格式，它定义了一套与框架无关的计算算子（Opset）和静态图描述。转换过程并非“翻译代码”，而是将训练框架中的计算图追踪并映射为ONNX的算子序列，同时把权重参数序列化为ONNX的TensorProto格式。

这一过程的关键在于算子映射：PyTorch的nn.Conv2d对应ONNX的Conv算子，TensorFlow的tf.nn.batch_normalization对应ONNX的BatchNormalization算子。但并非所有算子都能一一对应，例如PyTorch中的torch.where在旧版ONNX中可能被拆解为多个基础算子。因此，转换工具的核心工作就是构建一张“框架算子→ONNX算子”的映射表，并在遇到无法直接映射的算子时，尝试用多个ONNX基础算子组合等价替换。

4.2 方法流程：PyTorch导出与TensorFlow转换

4.2.1 PyTorch模型导出（torch.onnx.export）

PyTorch从1.2版本起提供了稳定的ONNX导出接口。核心函数是torch.onnx.export()，它通过追踪（Tracing）机制运行一次模型前向传播，记录实际执行过的Tensor操作，并将其转换为ONNX图。导出时需提供示例输入（dummy input），其形状和数据类型决定了ONNX图的输入签名。

import torch
import torchvision.models as models

# 1. 加载预训练模型并设为评估模式
model = models.resnet18(pretrained=True)
model.eval()

# 2. 构造示例输入（batch=1, 3通道, 224x224）
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出ONNX
torch.onnx.export(
    model,                      # 要导出的模型
    dummy_input,                # 示例输入
    "resnet18.onnx",            # 输出文件路径
    export_params=True,         # 导出权重参数
    opset_version=12,           # ONNX算子集版本
    do_constant_folding=True,   # 常量折叠优化
    input_names=["input"],      # 输入节点名
    output_names=["output"],    # 输出节点名
    dynamic_axes={              # 动态轴定义
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)
print("导出完成: resnet18.onnx")

关键参数说明：opset_version决定了可用的ONNX算子集合，版本越高支持的新算子越多，但兼容性可能下降；dynamic_axes允许指定哪些维度是动态的（如batch大小可变），这在部署时非常有用，但会增加图优化的复杂度。

4.2.2 TensorFlow模型转换（tf2onnx）

TensorFlow模型转换通常使用tf2onnx工具。它支持将SavedModel、GraphDef或Keras的.h5文件转换为ONNX。转换过程需要指定输入输出的张量名称，这些名称可以在TensorFlow的signature中查询。

# 命令行方式（推荐）
python -m tf2onnx.convert \
    --saved-model ./saved_model_dir \
    --output model.onnx \
    --opset 13 \
    --inputs serving_default_input:0 \
    --outputs StatefulPartitionedCall:0

# Python API方式
import tf2onnx
import tensorflow as tf

model = tf.keras.models.load_model("mobilenet.h5")
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
onnx_model = tf2onnx.convert.from_keras(model, input_signature=spec, opset=13)
with open("mobilenet.onnx", "wb") as f:
    f.write(onnx_model[0].SerializeToString())

4.3 案例：ResNet18图像分类模型转换与验证

下面以一个完整的ResNet18转换案例，展示从PyTorch到ONNX再到验证的完整流程。该案例模拟了实际部署中的标准操作。

# 步骤3：结构验证
import onnx
model_onnx = onnx.load("resnet18.onnx")
onnx.checker.check_model(model_onnx)
print("ONNX模型结构验证通过")

# 步骤5：数值比对
import onnxruntime as ort
import numpy as np

# 创建ONNX Runtime推理会话
sess = ort.InferenceSession("resnet18.onnx")
input_name = sess.get_inputs()[0].name

# 准备相同的输入数据
test_input = np.random.randn(1, 3, 224, 224).astype(np.float32)

# PyTorch推理
with torch.no_grad():
    torch_output = model(torch.from_numpy(test_input)).numpy()

# ONNX Runtime推理
onnx_output = sess.run(None, {input_name: test_input})[0]

# 计算最大绝对误差
max_error = np.max(np.abs(torch_output - onnx_output))
print(f"最大数值误差: {max_error:.6f}")
assert max_error < 1e-4, "数值误差超出阈值！"

4.4 常见误区与排障复盘

排障复盘：算子不支持错误。当遇到Unsupported operator: aten::unique之类的错误时，首先检查ONNX算子集版本是否过旧（升级opset_version），其次查看该算子是否有ONNX等价实现。若确实无法映射，可考虑在模型定义中替换该算子，例如用torch.unique的替代实现或自定义ONNX算子。

排障复盘：动态轴设置不当。如果部署时输入batch大小会变化，但导出时未设置dynamic_axes，ONNX Runtime会固定输入形状，导致推理报错。正确做法是在导出时明确声明动态维度，并在Java端使用OrtSession.InputOptions配置动态输入。