模型导出基础：从HuggingFace到ONNX与TorchScript

模型导出这件事，说白了就是把训练好的模型「换个格式」存下来。你可能觉得这有什么难的？不就是torch.save()吗？嗯，我刚开始也这么想，直到第一次在TensorRT-LLM上跑推理时，发现模型根本加载不进去……

今天咱们就聊聊三种最常见的导出方式：HuggingFace原生导出、ONNX导出、TorchScript导出。每种方式都有它的脾气，选对了事半功倍。

HuggingFace模型导出：最省心的方式

如果你用的是Transformers库，那导出模型其实就一行代码的事。我个人习惯用save_pretrained()方法，它会同时保存模型权重和配置文件。

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

model.save_pretrained("./llama2-7b-hf")
tokenizer.save_pretrained("./llama2-7b-hf")

这样导出的目录结构大概是这样的：

./llama2-7b-hf/
├── config.json          # 模型配置
├── pytorch_model.bin    # 权重文件（可能被分片）
├── tokenizer.json       # 分词器
├── tokenizer_config.json
└── generation_config.json
小提示：大模型通常会被自动分片成多个.bin文件，比如pytorch_model-00001-of-00003.bin。这是为了绕过单个文件的大小限制，TensorRT-LLM是支持直接加载这种分片格式的。

这种方式的优点很明显——简单、完整、生态好。但缺点呢？它保存的是PyTorch的动态图，TensorRT-LLM没法直接吃。你需要额外做一步转换。

ONNX导出：跨平台的「通用语言」

ONNX（Open Neural Network Exchange）就像模型界的普通话。不管你是PyTorch、TensorFlow还是其他框架训练的，都能转成ONNX格式。我在项目中遇到过好几次，甲方要求必须用ONNX部署，因为他们的推理框架只认这个。

导出ONNX的核心代码：

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("./llama2-7b-hf")
model.eval()

dummy_input = torch.randint(0, 32000, (1, 128))  # 模拟输入

torch.onnx.export(
    model,
    dummy_input,
    "llama2-7b.onnx",
    input_names=["input_ids"],
    output_names=["logits"],
    dynamic_axes={
        "input_ids": {0: "batch_size", 1: "sequence_length"},
        "logits": {0: "batch_size", 1: "sequence_length"}
    },
    opset_version=17
)
注意：ONNX导出时，dynamic_axes参数很关键。如果不设置，导出的模型就固定死了输入尺寸，推理时只能处理同样长度的序列。我曾经因为这个踩过坑，模型导出后只能处理128长度的输入，换成长文本直接报错。

ONNX的优势在于：

  • 跨框架兼容——PyTorch导出的ONNX，可以在TensorRT、OpenVINO上跑
  • 图优化——ONNX Runtime会自动做算子融合、常量折叠
  • 量化支持——可以导出INT8、FP16的ONNX模型

但缺点也很明显：

  • 导出过程可能遇到算子不兼容的问题
  • 动态形状支持不够完美
  • 对于大模型，ONNX文件体积可能比原始权重还大

TorchScript导出：PyTorch的亲儿子

TorchScript是PyTorch官方推出的序列化格式。它有两种导出方式：torch.jit.trace（追踪）和torch.jit.script（脚本）。

追踪模式（Trace）：

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("./llama2-7b-hf")
model.eval()

dummy_input = torch.randint(0, 32000, (1, 128))
traced_model = torch.jit.trace(model, dummy_input)
traced_model.save("llama2-7b_traced.pt")

脚本模式（Script）：

scripted_model = torch.jit.script(model)
scripted_model.save("llama2-7b_scripted.pt")
关键区别：Trace模式只记录实际执行过的路径，遇到if-else分支会出问题。Script模式会完整编译整个模型图，但要求代码必须符合TorchScript语法。我个人更推荐Trace模式，因为它简单、稳定，只要你的模型没有动态控制流，基本不会翻车。

TorchScript的好处是：

  • 与PyTorch无缝集成
  • 支持C++部署
  • 可以做图优化

但说实话，在TensorRT-LLM的部署流程里，TorchScript用得不多。因为TensorRT-LLM更倾向于直接吃HuggingFace格式或者ONNX格式。

三种格式对比：一张表说清楚

对比维度 HuggingFace ONNX TorchScript
导出难度 ⭐（最简单） ⭐⭐⭐（中等） ⭐⭐（较简单）
跨平台兼容性 低（仅PyTorch） 高（几乎所有框架） 中（PyTorch/C++）
TensorRT-LLM支持 直接支持 需要转换 不直接支持
动态形状支持 原生支持 需配置dynamic_axes Trace不支持，Script支持
图优化能力 无 强（ONNX Runtime） 中（TorchScript JIT）
文件体积 较小 较大 中等
调试难度 低 高（算子兼容性问题） 中（Script语法限制）

我的选择建议

如果你问我怎么选，我的答案是：看场景。

  • 做TensorRT-LLM部署：直接用HuggingFace格式最省事。TensorRT-LLM有专门的脚本把HF模型转成自己的格式，你不需要折腾ONNX或TorchScript。
  • 需要跨平台部署：选ONNX。比如你的模型要在NVIDIA GPU和Intel CPU上都跑，ONNX是唯一的选择。
  • 纯PyTorch生态：TorchScript够用了，特别是你要用LibTorch做C++推理的时候。
避坑指南：我曾经在一个项目里，为了追求「统一格式」，强行把所有模型都转成ONNX。结果有个模型里的自定义算子ONNX不支持，折腾了两天才找到替代方案。后来我学乖了——能用HF格式就用HF格式，别为了「统一」而「统一」。

知识体系总览

下面这张图帮你理清三种导出方式的关系和适用场景：

模型导出方式对比与选择 训练好的模型 HuggingFace格式 save_pretrained() ONNX格式 torch.onnx.export() TorchScript格式 torch.jit.trace/script ✅ 最简单 ✅ TensorRT-LLM原生支持 ❌ 仅PyTorch生态 ✅ 跨平台 ✅ 图优化强 ❌ 算子兼容问题 ✅ C++部署 ✅ 与PyTorch无缝 ❌ TensorRT-LLM不直接支持 💡 推荐：TensorRT-LLM部署 → 直接用HuggingFace格式 跨平台需求 → ONNX | 纯PyTorch C++ → TorchScript 三种格式各有优劣，根据实际部署场景选择最合适的

嗯，模型导出这块其实不难，关键是搞清楚每种格式的「脾气」。下一章咱们会深入TensorRT-LLM的模型转换流程，到时候你就知道为什么我推荐直接用HuggingFace格式了。

专注资料整理
← 上一章 📖 返回目录 下一章 →