跳转至内容

如何使用Ultralytics导出YOLO PyTorch

将PyTorch 部署到生产环境通常意味着需要针对每个目标平台分别使用不同的导出器： torch.onnx.export 对于ONNX， coremltools 对于 Apple 设备， onnx2tf 对于 TTensorFlow， pnnx 例如NCNN 等。每种工具都有其专属的 API、依赖项的特殊性以及输出规范。

Ultralytics 独立的导出工具，这些工具将多个后端整合到一个统一的界面中。您可以导出任何 torch.nn.Module，包括 timm 图像模型， 火炬视觉 分类器和检测器，或您自己的自定义架构，以 ONNX, TorchScript, OpenVINO, CoreML, NCNN, PaddlePaddle, MNN, ExecuTorch和 TensorFlow SavedModel 而无需分别学习每个后端。

为什么选择Ultralytics YOLO ？

  • 一个 API 支持 10 种格式：只需掌握一种调用规范，而非十几种。
  • 共享公用区域： 导出辅助程序位于 ultralytics.utils.export因此，一旦后端包安装完毕，您就可以在各种格式间保持相同的调用模式。
  • 与YOLO 采用相同的代码路径：所有Ultralytics YOLO 相同的辅助函数驱动。
  • 对于支持该功能的格式（OpenVINO、CoreML、MNN、NCNN），已内置FP16 和 INT8 量化功能。
  • 在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 以及一个示例输入tensor。MNN、TF SavedModel 和TF 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 图keras2pb()请参阅下方的详细要求.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 Graph 的导出都需要通过ONNX 中间步骤。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)

请务必致电 model.eval() 导出前

辍学， 批量归一化，以及其他仅用于训练的层在推理过程中表现不同。跳过 .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 对于动态输入形状， 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")

关于 分类 模型，将类名的列表传递给 classifier_names 向CoreML 添加一个分类头。

需要 coremltools>=9.0, torch>=1.11和 numpy<=2.3.5. Windows 系统不支持此功能。

BlobWriter not loaded 错误

coremltools>=9.0 在 macOS 和 Linux 系统上，Ships 支持Python .Python –3.13 版本。在更新的Python ，原生 C 扩展无法加载。请使用Python .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 图

承接上文的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() 检查 ncnn 和 pnnx 首次使用时。

导出到 MNN

MNN 导出需要ONNX 作为输入。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.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 ：

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

NVIDIA 不支持此功能。

导出到 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.forward(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 导出版本的容差范围更宽。请使用真实数据而非随机张量进行验证。

对于其他运行时，输入tensor 可能会有所不同。例如，OpenVINO 使用模型的前向参数名称（通常为 x （对于通用模型），而 torch2onnx 默认设置为 "images".

已知限制

  • 多输入支持情况不一: torch2onnx 和 torch2openvino 对于具有多个输入的模型，接受一个示例张量的元组或列表。 torch2torchscript, torch2coreml, torch2ncnn, torch2paddle和 torch2executorch 假设有一个输入tensor。
  • ExecuTorch 需要 flatc: ExecuTorch 运行时需要 FlatBuffers 编译器。请使用以下命令进行安装： brew install flatbuffers 在 macOS 上或 apt install flatbuffers-compiler 在 Ubuntu 上。
  • Ultralytics通过Ultralytics进行推断: 导出的YOLO 无法通过 YOLO() 用于推理。请为每种格式使用相应的原生运行时（ONNX Runtime, OpenVINO，等等）。
  • YOLO 的格式：Axelera和Sony IMX500的导出文件需要YOLO模型属性，且不适用于通用模型。
  • 特定平台的格式: TensorRT 需要NVIDIA GPU。 RKNN 需要 rknn-toolkit2 SDK（仅限 Linux）。 Edge TPU 需要 edgetpu_compiler 二进制文件（仅限 Linux）。

常见问题

使用Ultralytics 可以导出哪些模型？

任何 torch.nn.Module。这包括来自 timm、torchvision 或任何自定义PyTorch 模型。该模型必须处于评估模式（model.eval()) 导出前。OpenVINO ONNX OpenVINO 为多输入模型提供示例张量的元组。

哪些导出格式GPU即可使用？

所有受支持的格式（TorchScript、ONNX、OpenVINO、CoreML、TF SavedModel、TF 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 通过 torch2coreml. 对于图像分类模型，请将类名的列表传递给 classifier_names 在分类器中进行训练。请在 macOS 或 Linux 上运行导出操作。Windows 不CoreML 。请参阅 CoreML 有关iOS 详细信息。

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

是的，支持多种格式。通过 half=True 适用于 FP16 或 int8=True 在导出至OpenVINO、CoreML、MNN 或NCNN 时，请使用 INT8。OpenVINO 中的 INT8OpenVINO 要求 calibration_dataset 参数用于 训练后量化. 有关各格式的量化权衡，请参阅相应格式的集成页面。



📅 创建于 0 天前 ✏️ 更新于 0 天前
raimbekovm

评论