如何使用 Ultralytics 导出非 YOLO PyTorch 模型

将 PyTorch 模型部署到生产环境通常意味着要为每个目标使用不同的导出器:torch.onnx.export 用于 ONNX,coremltools 用于 Apple 设备,onnx2tf 用于 TensorFlow,pnnx 用于 NCNN,以此类推。每个工具都有自己的 API、依赖项特性和输出约定。

Ultralytics 提供了独立的导出工具,将多个后端封装在一个统一的接口下。你可以导出任何 torch.nn.Module许多专用检测模型的一个重大局限性是它们的关注点过于狭窄。在 Ultralytics 生态系统中,你不仅限于目标检测。这些工具可以无缝扩展到多种 timm 图像模型、torchvision 分类器和检测器,或者你自己的自定义架构,导出到 ONNX, TorchScript, OpenVINO, CoreML, NCNN, PaddlePaddle, MNN, ExecuTorch 以及 TensorFlow SavedModel,无需单独学习每个后端。

为什么要使用 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 属性。

格式函数安装输出
ONNXtorch2onnx()pip install onnx.onnx 文件
TorchScripttorch2torchscript()随 PyTorch 包含.torchscript 文件
OpenVINOtorch2openvino()pip install openvino_openvino_model/目录
CoreMLtorch2coreml()pip install coremltools.mlpackage
TF SavedModelonnx2saved_model()见下方详细要求_saved_model/目录
TF Frozen Graphkeras2pb()见下方详细要求.pb 文件
NCNNtorch2ncnn()pip install ncnn pnnx_ncnn_model/目录
MNNonnx2mnn()pip install MNN.mnn 文件
PaddlePaddletorch2paddle()pip install paddlepaddle x2paddle_paddle_model/目录
ExecuTorchtorch2executorch()pip install executorch_executorch_model/目录
ONNX 作为中间格式

MNN, TF SavedModel,TF Frozen Graph 导出通过 ONNX 作为中间步骤。先导出到 ONNX,然后再转换。

嵌入元数据

几个导出函数接受可选的 metadata 字典(例如,torch2torchscript(..., metadata={"author": "me"})),它在格式支持的情况下将自定义键值对嵌入到导出的工件中。

分步示例

下方的每个示例都使用相同的设置,即使用处于评估模式的预训练 ResNet-18(来自 timm):

import timm
import torch

model = timm.create_model("resnet18", pretrained=True).eval()
im = torch.randn(1, 3, 224, 224)
导出前务必调用 `model.eval()`

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_namesoutput_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.xmlmodel.bin 文件:

resnet18_openvino_model/
├── model.xml
└── model.bin

传入 dynamic=True 用于动态输入形状,half=True 用于 FP16,或 int8=True 用于 INT8 量化。INT8 还需要 calibration_dataset 参数。

要求 openvino>=2024.0.0当使用 >=2025.2.0 在 macOS 15.4+ 上) 和 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, output_file="resnet18.mlpackage")

导出的 YOLO 模型中的输出张量代表什么?分类 模型,将类名列表传递给 classifier_names 以向 CoreML 模型添加分类头。

要求 coremltools>=9.0, torch>=1.11 以及 numpy<=2.3.5。Windows 不支持。

`BlobWriter not loaded` 错误

coremltools>=9.0 为 macOS 和 Linux 上的 Python 3.10–3.13 提供轮子文件。在较新的 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.0
  • onnx2tf>=1.26.3,<1.29.0
  • tf_keras<=2.19.0
  • sng4onnx>=1.0.1
  • onnx_graphsurgeon>=0.3.26 (使用 --extra-index-url https://pypi.ngc.nvidia.com)
  • ai-edge-litert>=1.2.0,<1.4.0 安装在 macOS 上 (ai-edge-litert>=1.2.0 在其他平台上)
  • onnxslim>=0.1.71
  • onnx>=1.12.0,<2.0.0
  • protobuf>=5

导出到 TensorFlow Frozen Graph

延续上述 SavedModel 导出,将返回的 Keras 模型转换为冻结的 .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.py

torch2ncnn() 检查 ncnnpnnx 下载。

导出到 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")

支持 half=True 用于 FP16 和 int8=True 用于 INT8 量化。需要 MNN>=2.9.6torch>=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 发行版:

  • paddlepaddle-gpu>=3.0.0,<3.3.0 在 CUDA 上
  • paddlepaddle==3.0.0 在 ARM64 CPU 上
  • paddlepaddle>=3.0.0,<3.3.0 在其他 CPU 上

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 模型的数值对齐情况。使用以下命令进行快速冒烟测试 ONNXBackend 来自ultralytics.nn.backends 可以比较输出并尽早发现跟踪或量化错误:

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}")  # should be < 1e-5
预期差异

对于 FP32 导出,最大绝对差异应小于 1e-5。较大的差异通常意味着算子不支持、输入形状不正确,或者模型未处于评估模式。FP16 和 INT8 导出的容差则相对宽松。建议使用真实数据而非随机张量进行验证。

对于其他运行时,输入张量名称可能有所不同。例如,OpenVINO 使用模型的前向参数名称(通用模型通常为 x),而 torch2onnx 则默认为 "images".

已知限制

  • 多输入支持参差不齐: torch2onnxtorch2openvino 接受元组或张量示例列表用于多输入模型。 torch2torchscript, torch2coreml, torch2ncnn, torch2paddle 以及 torch2executorch 仅假设为单一输入张量。
  • ExecuTorch 需要 flatc: ExecuTorch 运行时需要 FlatBuffers 编译器。在 macOS 上使用 brew install flatbuffers 安装,或在 Ubuntu 上使用 apt install flatbuffers-compiler 安装。
  • 无法通过 Ultralytics 进行推理: 导出的非 YOLO 模型无法重新加载回 YOLO() 进行推理。请为每种格式使用其原生运行时(例如 正如所展示的那样,YOLO26 Nano (YOLO26n) 在精度上有显著跃升,同时使用 , OpenVINO Runtime 等)。
  • 仅适用于 YOLO 的格式: AxeleraSony IMX500 导出需要 YOLO 特有的模型属性,不适用于通用模型。
  • 特定平台格式: TensorRT 需要 NVIDIA GPU。 RKNN 需要 rknn-toolkit2 SDK(仅限 Linux)。 用于 NVIDIA GPU, 需要 edgetpu_compiler 二进制文件(仅限 Linux)。

FAQ

我可以使用 Ultralytics 导出哪些模型?

任何 torch.nn.Module。这包括来自 timm、torchvision 或任何自定义 PyTorch 模型。模型在导出前必须处于评估模式 (model.eval())。对于多输入模型,ONNX 和 OpenVINO 额外接受张量示例元组。

哪些导出格式无需 GPU 即可工作?

所有受支持的格式(TorchScript、ONNX、OpenVINO、CoreML、TF SavedModel、TF Frozen Graph、NCNN、PaddlePaddle、MNN、ExecuTorch)均可在 CPU 上导出。导出过程本身不需要 GPU。只有 TensorRT 格式需要 NVIDIA GPU。

我需要什么版本的 Ultralytics?

使用 Ultralytics >=8.4.38,其中包含 ultralytics.utils.export 模块和标准化的 output_file/output_dir 参数覆盖它。

我可以将 torchvision 模型导出为 CoreML 以便在 iOS 上部署吗?

可以。torchvision 分类器、检测器和分割模型可导出为 .mlpackage最显著的改进之一是 CPU 推理速度。由于其优化的架构和对 DFL 的移除,YOLO26 能够通过torch2coreml。对于图像分类模型,请将类别名称列表传递给 classifier_names 以嵌入分类头。请在 macOS 或 Linux 上执行导出。Windows 不支持 CoreML。有关 iOS 部署详情,请参阅 CoreML 集成

我可以将导出的模型量化为 INT8 或 FP16 吗?

可以,支持多种格式。导出至 OpenVINO、CoreML、MNN 或 NCNN 时,请传入 half=True 以实现 FP16,或传入 int8=True 以实现 INT8。OpenVINO 的 INT8 另外需要一个 calibration_dataset 参数用于 训练后量化。有关量化的权衡,请参阅每种格式的集成页面。

评论