如何通过 Ultralytics 导出非 YOLO PyTorch 模型#
Ultralytics 在 ultralytics.utils.export 下提供独立的导出工具,将多个后端封装在一个统一的接口后方。你可以将任何 torch.nn.Module,包括 timm 图像模型、torchvision 分类器和检测器,或你自己的自定义架构,导出为 ONNX、TorchScript、OpenVINO、CoreML、NCNN、PaddlePaddle、MNN、ExecuTorch 和 TensorFlow SavedModel,而无需单独学习每个后端。
将 PyTorch 模型部署到生产环境通常意味着针对每个目标都需要应付不同的导出器:用于 ONNX 的 torch.onnx.export、用于 Apple 设备的 coremltools、用于 TensorFlow 的 onnx2tf、用于 NCNN 的 pnnx,等等。每个工具都有其自己的 API、依赖项怪癖和输出约定。这些工具将它们简化为了单一的调用模式。
为什么要使用 Ultralytics 进行非 YOLO 模型导出?#
- 跨 10 种格式的单一 API: 学习一套调用约定即可,无需掌握十几种。
- 共享的实用工具表面: 导出辅助工具位于
ultralytics.utils.export下,因此一旦安装了后端包,你就可以在各种格式之间保持相同的调用模式。 - 与 YOLO 导出相同的代码路径: 同样的辅助函数驱动着每一次 Ultralytics YOLO 的导出。
- 内置 FP16 和 INT8 量化: 支持这些功能的格式(OpenVINO、CoreML、MNN、NCNN)均已内置此功能。
- 在 CPU 上运行: 导出步骤本身不需要 GPU,因此你可以在任何笔记本电脑上本地运行。
快速入门#
最快的路径是通过两行代码将模型导出到 ONNX,其中不包含 YOLO 代码,并且除了 pip install ultralytics onnx timm 之外无需其他设置:
import timm
import torch
from ultralytics.utils.export import torch2onnx
model = timm.create_model("resnet18", pretrained=True).eval()
torch2onnx(model, torch.randn(1, 3, 224, 224), output_file="resnet18.onnx")支持的导出格式#
torch2* 函数接受标准的 torch.nn.Module 和示例输入张量。MNN、TF SavedModel 和 TF Frozen Graph 会通过中间的 ONNX 或 Keras 工件。在任何情况下都不需要 YOLO 特定的属性。
| 格式 | 函数 | 安装 | 输出 |
|---|---|---|---|
| ONNX | torch2onnx() | pip install onnx | .onnx 文件 |
| TorchScript | torch2torchscript() | 包含在 PyTorch 中 | .torchscript 文件 |
| OpenVINO | torch2openvino() | pip install openvino | _openvino_model/ 目录 |
| CoreML | torch2coreml() | pip install coremltools | .mlpackage |
| TF SavedModel | onnx2saved_model() | 参见下方的详细要求 | _saved_model/ 目录 |
| TF Frozen Graph | keras2pb() | 参见下方的详细要求 | .pb 文件 |
| NCNN | torch2ncnn() | pip install ncnn pnnx | _ncnn_model/ 目录 |
| MNN | onnx2mnn() | pip install MNN | .mnn 文件 |
| PaddlePaddle | torch2paddle() | pip install paddlepaddle x2paddle | _paddle_model/ 目录 |
| ExecuTorch | torch2executorch() | pip install executorch | _executorch_model/ 目录 |
MNN、TF SavedModel 和 TF Frozen Graph 导出将 ONNX 作为中间步骤。请先导出为 ONNX,然后再进行转换。
几个导出函数接受一个可选的 metadata 字典(例如 torch2torchscript(..., metadata={"author": "me"})),该字典在格式支持的情况下将自定义键值对嵌入到导出的工件中。
分步示例#
下方的每个示例都使用相同的设置:一个来自 timm 的预训练 ResNet-18,并处于评估模式:
import timm
import torch
model = timm.create_model("resnet18", pretrained=True).eval()
im = torch.randn(1, 3, 224, 224)Dropout、批归一化 (batch normalization) 和其他仅限训练的层在推理过程中的行为有所不同。跳过 .eval() 会导致导出的模型输出错误。
导出到 ONNX#
from ultralytics.utils.export import torch2onnx
torch2onnx(model, im, output_file="resnet18.onnx")对于动态批大小,请传入 dynamic 字典:
torch2onnx(model, im, output_file="resnet18_dyn.onnx", dynamic={"images": {0: "batch_size"}})默认的 opset 是 14,默认的输入名称是 "images"。可以使用 opset、input_names 或 output_names 参数进行覆盖。
导出到 TorchScript#
无需额外依赖。底层使用 torch.jit.trace。
from ultralytics.utils.export import torch2torchscript
torch2torchscript(model, im, output_file="resnet18.torchscript")导出到 OpenVINO#
from ultralytics.utils.export import torch2openvino
ov_model = torch2openvino(model, im, output_dir="resnet18_openvino_model")该目录包含一对固定名称的 model.xml 和 model.bin:
resnet18_openvino_model/
├── model.xml
└── model.bin为动态输入形状传入 dynamic=True,为 FP16 传入 quantize=16,或为 INT8 量化传入 quantize=8。INT8 额外需要 calibration_dataset 参数。
需要 openvino>=2024.0.0(或在 macOS 15.4+ 上为 >=2025.2.0)以及 torch>=2.1。
导出到 CoreML#
import coremltools as ct
from ultralytics.utils.export import torch2coreml
inputs = [ct.TensorType("input", shape=(1, 3, 224, 224))]
ct_model = torch2coreml(model, inputs, im, classifier_names=None, output_file="resnet18.mlpackage")对于分类模型,将类别名称列表传递给 classifier_names,以便向 CoreML 模型添加分类头。
需要 coremltools>=9.0、torch>=1.11 和 numpy<=2.3.5。Windows 上不支持。
coremltools>=9.0 提供适用于 macOS 和 Linux 上 Python 3.10–3.13 的 wheel 包。在较新的 Python 版本上,原生 C 扩展会加载失败。请使用 Python 3.10–3.13 进行 CoreML 导出。
导出到 TensorFlow SavedModel#
TF SavedModel 导出通过 ONNX 作为中间步骤:
from ultralytics.utils.export import onnx2saved_model, torch2onnx
torch2onnx(model, im, output_file="resnet18.onnx")
keras_model = onnx2saved_model("resnet18.onnx", output_dir="resnet18_saved_model")该函数返回一个 Keras 模型,并会在输出目录中生成 TFLite 文件(.tflite):
resnet18_saved_model/
├── saved_model.pb
├── variables/
├── resnet18_float32.tflite
├── resnet18_float16.tflite
└── resnet18_int8.tflite要求:
tensorflow>=2.0.0,<=2.19.0onnx2tf>=1.26.3,<1.29.0tf_keras<=2.19.0sng4onnx>=1.0.1onnx_graphsurgeon>=0.3.26- macOS 上的
ai-edge-litert>=1.2.0,<1.4.0(其他平台上为ai-edge-litert>=1.2.0) onnxslim>=0.1.82onnx>=1.12.0,<2.0.0protobuf>=5
导出到 TensorFlow Frozen Graph#
接上面的 SavedModel 导出,将返回的 keras_model 转换为冻结的 .pb 图:
from pathlib import Path
from ultralytics.utils.export import keras2pb
keras2pb(keras_model, output_file=Path("resnet18_saved_model/resnet18.pb"))导出到 NCNN#
from ultralytics.utils.export import torch2ncnn
torch2ncnn(model, im, output_dir="resnet18_ncnn_model")目录包含固定名称的 param 和 bin 文件,以及一个 Python 包装器:
resnet18_ncnn_model/
├── model.ncnn.param
├── model.ncnn.bin
└── model_ncnn.pytorch2ncnn() 在首次使用时会检查 ncnn 和 pnnx。
导出到 MNN#
MNN 导出需要一个 ONNX 文件作为输入。请先导出到 ONNX,然后再进行转换:
from ultralytics.utils.export import onnx2mnn, torch2onnx
torch2onnx(model, im, output_file="resnet18.onnx")
onnx2mnn("resnet18.onnx", output_file="resnet18.mnn")支持用于 FP16 的 quantize=16 和用于 INT8 量化的 quantize=8。需要 MNN>=2.9.6 和 torch>=1.10。
导出到 PaddlePaddle#
from ultralytics.utils.export import torch2paddle
torch2paddle(model, im, output_dir="resnet18_paddle_model")目录包含 PaddlePaddle 模型和参数文件:
resnet18_paddle_model/
├── model.pdmodel
└── model.pdiparams需要 x2paddle 以及适用于你平台的正确 PaddlePaddle 发行版:
- CUDA 上的
paddlepaddle-gpu>=3.0.0,<3.3.0 - ARM64 CPU 上的
paddlepaddle==3.0.0 - 其他 CPU 上的
paddlepaddle>=3.0.0,<3.3.0
不支持 NVIDIA Jetson。
导出为 ExecuTorch#
from ultralytics.utils.export import torch2executorch
torch2executorch(model, im, output_dir="resnet18_executorch_model")导出的 .pte 文件保存在输出目录中:
resnet18_executorch_model/
└── model.pte需要 torch>=2.9.0 和匹配的 ExecuTorch 运行时(pip install executorch)。有关运行时用法,请参阅 ExecuTorch 集成。
验证你导出的模型#
导出后,在发布前验证与原始 PyTorch 模型的数值一致性。使用来自 ultralytics.nn.backends 的 ONNXBackend 进行快速冒烟测试,可以比较输出并尽早标记出追踪或量化错误:
import numpy as np
import timm
import torch
from ultralytics.nn.backends import ONNXBackend
model = timm.create_model("resnet18", pretrained=True).eval()
im = torch.randn(1, 3, 224, 224)
with torch.no_grad():
pytorch_output = model(im).numpy()
onnx_model = ONNXBackend("resnet18.onnx", device=torch.device("cpu"))
onnx_output = onnx_model(im)[0]
diff = np.abs(pytorch_output - onnx_output).max()
print(f"Max difference: {diff:.6f}") # typically ~1e-5, well under 1e-4 for FP32对于 FP32 导出,最大绝对差通常约为 1e-5,并且应远低于 1e-4。较大的差异表明存在不支持的操作、错误的输入形状或模型不在评估模式下。FP16 和 INT8 导出的公差较宽松。请使用真实数据而不是随机张量进行验证。
对于其他运行时,输入张量名称可能会有所不同。例如,OpenVINO 使用模型的向前参数名称(通用模型通常为 x),而 torch2onnx 默认值为 "images"。
已知限制#
- 对多输入的支持不均衡:对于具有多个输入的模型,
torch2onnx和torch2openvino接受示例张量的元组或列表。torch2torchscript、torch2coreml、torch2ncnn、torch2paddle和torch2executorch假定为单个输入张量。 - ExecuTorch 需要
flatc:ExecuTorch 运行时需要 FlatBuffers 编译器。在 macOS 上使用brew install flatbuffers安装,或在 Ubuntu 上使用apt install flatbuffers-compiler安装。 - 无法通过 Ultralytics 进行推理:导出的非 YOLO 模型无法通过
YOLO()重新加载以进行推理。请为每种格式使用原生运行时(ONNX Runtime、OpenVINO Runtime 等)。 - 仅限 YOLO 的格式:Axelera 和 Sony IMX500 导出需要 YOLO 特定的模型属性,不适用于通用模型。
- 平台特定格式:TensorRT 需要 NVIDIA GPU。RKNN 需要
rknn-toolkit2SDK(仅限 Linux)。Edge TPU 需要edgetpu_compiler二进制文件(仅限 Linux)。
结论#
这些实用工具通过一个统一的 API,将任何 PyTorch 模型从普通的 torch.nn.Module 转换为可用于部署的 ONNX、OpenVINO、CoreML、TensorFlow 或移动运行时工件。选择与你的目标硬件匹配的格式,针对原始模型验证数值一致性,然后按照匹配的集成指南进行运行时特定的部署步骤。
常见问题解答#
任何
torch.nn.Module。这包括来自 timm、torchvision 或任何自定义 PyTorch 模型的模型。模型在导出前必须处于评估模式(model.eval())。ONNX 和 OpenVINO 额外接受用于多输入模型的示例张量元组。所有支持的格式(TorchScript, ONNX, OpenVINO, CoreML, TF SavedModel, TF Frozen Graph, NCNN, PaddlePaddle, MNN, ExecuTorch)都可以在 CPU 上导出。导出过程本身不需要 GPU。TensorRT 是唯一需要 NVIDIA GPU 的格式。
使用 Ultralytics
>=8.4.38,它包含ultralytics.utils.export模块以及标准化的output_file/output_dir参数。可以。torchvision 分类器、检测器和分割模型通过
torch2coreml导出为.mlpackage。对于图像分类模型,请将类别名称列表传递给classifier_names以内嵌分类头。请在 macOS 或 Linux 上运行导出。Windows 不支持 CoreML。有关 iOS 部署的详细信息,请参阅CoreML 集成。可以,适用于几种格式。在导出到 OpenVINO、CoreML、MNN 或 NCNN 时,传递
quantize=16用于 FP16 或传递quantize=8用于 INT8。OpenVINO 中的 INT8 另外需要一个calibration_dataset参数用于训练后量化。有关量化的权衡,请参阅各格式的集成页面。在相同的输入上运行原始 PyTorch 模型和导出的模型,然后比较输出。使用匹配的后端(例如用于 ONNX 的
ONNXBackend)加载导出的文件,并检查最大绝对差。对于 FP32 导出,它通常约为1e-5,并且应远低于1e-4;较大的差距表明存在不支持的操作、错误的输入形状或模型不在评估模式下。有关可运行的示例,请参阅验证你的导出模型。