3、torch.onnx.export核心参数：model、args、f、export_params等基础参数详解

在ONNX导出链路中，torch.onnx.export是连接PyTorch动态图世界与ONNX静态图世界的核心桥梁。该函数拥有超过20个可配置参数，但真正决定导出成败与质量的关键，往往集中在少数几个基础参数上。理解model、args、f、export_params这四个参数，是掌握ONNX导出的第一步，也是后续调优opset与dynamic_axes的前提。本章将逐一拆解这些参数的本质、使用方式与常见陷阱。

上图展示了四个核心参数在导出流程中的位置与作用：model提供计算图定义，args提供输入样例以触发实际执行，f指定输出目标，export_params控制权重存储方式。四者相互配合，共同决定最终ONNX模型的结构与行为。

3.1 model参数：导出源模型的正确姿势

核心原理

model参数接收一个torch.nn.Module实例，它是导出的源头。PyTorch通过追踪（tracing）机制，在给定输入下实际执行模型的前向传播，记录所有张量操作，并将其转换为ONNX计算图。这意味着model不仅需要是一个合法的Module，还必须处于可执行状态。

操作流程

1. 实例化模型并加载权重：model = MyModel(); model.load_state_dict(torch.load('weights.pth'))
2. 调用model.eval()切换到评估模式
3. 将模型移动到目标设备（CPU或GPU），但推荐在CPU上导出以避免设备相关算子差异
4. 将模型作为第一个参数传入torch.onnx.export

案例拆解：一个包含Dropout的模型导出

假设我们有一个包含Dropout层的简单分类模型：

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(256, 10)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.dropout(x)  # 训练时随机置零，评估时恒等映射
        return self.fc2(x)

model = SimpleNet()
# 错误示范：忘记调用eval()
torch.onnx.export(model, torch.randn(1, 784), "wrong.onnx")
# 正确做法
model.eval()
torch.onnx.export(model, torch.randn(1, 784), "correct.onnx")

在错误示范中，Dropout层在追踪时被记录为随机置零操作，导致导出的ONNX模型每次推理结果不同。正确调用eval()后，Dropout被识别为恒等映射，模型行为确定。

常见误区

• 误区1：认为model.cuda()后必须在GPU上导出。实际上，在CPU上导出更安全，因为某些GPU专属算子（如cudnn.batch_norm）可能无法映射到ONNX。
• 误区2：忽略model.train()与model.eval()的状态切换。导出前必须确认处于eval()状态。
• 误区3：模型包含不可追踪的控制流（如Python的if依赖张量值）。追踪机制只会记录实际执行的分支，导致导出的图不完整。

3.2 args参数：输入样例的精确控制

核心原理

args参数为模型提供一组示例输入，PyTorch通过执行这些输入来触发前向传播并记录计算图。这组输入不仅决定了导出的计算图结构，还固定了输入张量的shape、dtype和设备类型。如果后续部署时输入shape与导出时不一致，且未配置dynamic_axes，推理将直接报错。

输入格式规范

案例拆解：多输入模型的导出

考虑一个双输入模型，输入为图像和文本特征：

class MultiModal(nn.Module):
    def forward(self, image, text_feat):
        fused = image + text_feat
        return fused

model = MultiModal().eval()
image_input = torch.randn(1, 3, 224, 224)
text_input = torch.randn(1, 512)
torch.onnx.export(model, (image_input, text_input), "multi.onnx")

这里args必须使用元组(image_input, text_input)，顺序与forward定义一致。若传成[image_input, text_input]列表，PyTorch也能接受，但官方推荐使用元组以避免歧义。

常见误区

• 误区1：args中传入torch.Tensor的子类（如nn.Parameter），可能导致追踪异常。
• 误区2：args的batch size设置过大，导致导出时内存占用过高。建议使用batch size=1进行导出。
• 误区3：忽略输入张量的requires_grad属性。导出时应确保输入requires_grad=False，否则计算图中会包含梯度节点。

3.3 f参数：输出路径与文件管理

核心原理

f参数指定ONNX模型的输出位置，可以是文件路径字符串或文件对象。当export_params=True时，权重内嵌于单个.onnx文件中；当export_params=False时，权重以外部数据形式存储在同名.onnx_data文件中。

路径处理规则

• 支持路径与相对路径：f="./models/model.onnx"
• 支持文件对象：with open("model.onnx", "wb") as fp: torch.onnx.export(model, args, fp)
• 文件扩展名建议使用.onnx，但非强制
• 若目录不存在，PyTorch不会自动创建，需提前os.makedirs

案例拆解：外部权重导出与文件管理

import os
os.makedirs("export", exist_ok=True)

# 方式1：权重内嵌（单文件）
torch.onnx.export(model, args, "export/model_embedded.onnx", export_params=True)

# 方式2：权重外部化（双文件）
torch.onnx.export(model, args, "export/model_external.onnx", export_params=False)
# 生成 model_external.onnx + model_external.onnx_data

外部权重模式适用于模型超过2GB的场景（ONNX协议限制），或需要频繁更新权重而保持图结构不变的场景。但部署时需要同时分发两个文件，并确保相对路径正确。

常见误区

• 误区1：使用f传入文件对象后，忘记关闭文件句柄，导致资源泄漏。
• 误区2：在export_params=False时，误以为权重会嵌入到.onnx文件中，导致部署时缺少数据文件。
• 误区3：路径中包含中文字符或空格，某些ONNX Runtime版本可能解析异常。

3.4 export_params参数：权重存储策略

核心原理

export_params是一个布尔参数，控制模型权重是否作为ONNX图的初始器（initializer）内嵌到模型文件中。当为True时，所有权重以TensorProto形式存储在.onnx文件内部；为False时，权重被外部化为独立文件，ONNX图中仅保留引用关系。

两种模式对比

案例拆解：大模型导出策略选择

假设我们有一个1.8GB的BERT模型，需要部署到生产环境：

# 方案A：内嵌权重（推荐）
torch.onnx.export(model, args, "bert.onnx", export_params=True)
# 结果：bert.onnx 约1.8GB，单文件部署

# 方案B：外部权重
torch.onnx.export(model, args, "bert_ext.onnx", export_params=False)
# 结果：bert_ext.onnx 约2MB + bert_ext.onnx_data 约1.8GB

方案A适合容器化部署，只需挂载一个文件；方案B适合需要热更新权重的场景，例如模型微调后仅替换.onnx_data文件。但需注意，ONNX Runtime加载外部权重时，会先读取.onnx文件中的路径信息，再加载对应数据文件，因此两个文件必须保持相对路径一致。

常见误区

• 误区1：认为export_params=False可以减小总存储占用。实际上，权重数据总量不变，只是存储位置不同。
• 误区2：在export_params=False时，使用onnx.load()加载模型后直接保存，可能导致外部权重丢失。需使用onnx.save_model(..., save_as_external_data=True)重新保存。
• 误区3：混淆export_params与do_constant_folding。前者控制权重存储位置，后者控制常量折叠优化，两者独立。

3.5 参数组合实战：一个完整的导出流程

综合以上四个参数，一个规范的导出流程如下：

import torch
import torchvision.models as models

# 1. 准备模型
model = models.resnet18(pretrained=True)
model.eval()

# 2. 准备输入样例
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 执行导出
torch.onnx.export(
    model,                    # model参数
    dummy_input,              # args参数
    "resnet18.onnx",          # f参数
    export_params=True,       # 权重内嵌
    opset_version=11,         # 后续章节详解
    do_constant_folding=True, # 常量折叠优化
    input_names=["input"],    # 后续章节详解
    output_names=["output"]   # 后续章节详解
)

# 4. 验证导出结果
import onnx
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)
print(f"模型大小: {onnx_model.ByteSize() / 1024 / 1024:.2f} MB")