Elgato_dark/ultralytics
1
0
Fork 0
mirror of https://github.com/ultralytics/ultralytics synced 2026-04-21 14:07:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 8

refactor: split Exporter export methods into per-format utility modules (#23914)

Browse source 
Signed-off-by: Onuralp SEZER <onuralp@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Lakshantha Dissanayake <lakshantha@ultralytics.com>
Co-authored-by: Jing Qiu <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: Laughing-q <1185102784@qq.com>
This commit is contained in:
Onuralp SEZER 2026-03-27 13:38:56 +03:00 committed by GitHub
parent 19bbfc4d4f
commit b73dce8813
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
23 changed files with 988 additions and 456 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
docs/en/reference/engine/exporter.md
View file

@ -15,10 +15,6 @@ keywords: YOLOv8, export formats, ONNX, TensorRT, CoreML, machine learning model
<br><br><hr><br>
## ::: ultralytics.engine.exporter.IOSDetectModel
<br><br><hr><br>
## ::: ultralytics.engine.exporter.NMSModel
<br><br><hr><br>
@ -27,10 +23,6 @@ keywords: YOLOv8, export formats, ONNX, TensorRT, CoreML, machine learning model
<br><br><hr><br>
## ::: ultralytics.engine.exporter.best_onnx_opset
<br><br><hr><br>
## ::: ultralytics.engine.exporter.validate_args
<br><br><hr><br>

16
docs/en/reference/utils/export/axelera.md Normal file
View file

@ -0,0 +1,16 @@
---
description: Axelera export utilities for converting ONNX models to Axelera format for deployment on Metis AI processors. Supports INT8 quantization and optimized inference for computer vision workloads on Axelera hardware.
keywords: Ultralytics, Axelera, model export, ONNX to Axelera, Metis AI processor, INT8 quantization, computer vision, AI inference, edge AI, hardware acceleration
---
# Reference for `ultralytics/utils/export/axelera.py`
!!! success "Improvements"
This page is sourced from [https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/axelera.py](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/axelera.py). Have an improvement or example to add? Open a [Pull Request](https://docs.ultralytics.com/help/contributing/) — thank you! 🙏
<br>
## ::: ultralytics.utils.export.axelera.onnx2axelera
<br><br>

24
docs/en/reference/utils/export/coreml.md Normal file
View file

@ -0,0 +1,24 @@
---
description: CoreML export utilities for converting PyTorch YOLO models to CoreML format for Apple devices. Supports iOS deployment with FP16/INT8 quantization, NMS pipeline integration, and optimized inference on Apple Silicon and mobile devices.
keywords: Ultralytics, CoreML, model export, PyTorch to CoreML, Apple iOS, macOS, Apple Silicon, INT8 quantization, FP16, NMS pipeline, mobile deployment, on-device inference, MLProgram, neural network
---
# Reference for `ultralytics/utils/export/coreml.py`
!!! success "Improvements"
This page is sourced from [https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/coreml.py](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/coreml.py). Have an improvement or example to add? Open a [Pull Request](https://docs.ultralytics.com/help/contributing/) — thank you! 🙏
<br>
## ::: ultralytics.utils.export.coreml.IOSDetectModel
<br><br><hr><br>
## ::: ultralytics.utils.export.coreml.pipeline_coreml
<br><br><hr><br>
## ::: ultralytics.utils.export.coreml.torch2coreml
<br><br>

4
docs/en/reference/utils/export/engine.md
View file

@ -11,6 +11,10 @@ keywords: Ultralytics, TensorRT export, ONNX export, PyTorch to ONNX, quantizati
<br>
## ::: ultralytics.utils.export.engine.best_onnx_opset
<br><br><hr><br>
## ::: ultralytics.utils.export.engine.torch2onnx
<br><br><hr><br>

16
docs/en/reference/utils/export/mnn.md Normal file
View file

@ -0,0 +1,16 @@
---
description: MNN export utilities for converting ONNX models to MNN format for efficient inference on mobile and embedded devices. Supports FP16 and INT8 weight quantization for optimized deployment using Alibaba's MNN framework.
keywords: Ultralytics, MNN, model export, ONNX to MNN, Alibaba MNN, mobile deployment, embedded systems, FP16, INT8 quantization, lightweight inference, edge deployment
---
# Reference for `ultralytics/utils/export/mnn.py`
!!! success "Improvements"
This page is sourced from [https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/mnn.py](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/mnn.py). Have an improvement or example to add? Open a [Pull Request](https://docs.ultralytics.com/help/contributing/) — thank you! 🙏
<br>
## ::: ultralytics.utils.export.mnn.onnx2mnn
<br><br>

16
docs/en/reference/utils/export/ncnn.md Normal file
View file

@ -0,0 +1,16 @@
---
description: NCNN export utilities for converting PyTorch YOLO models to NCNN format using PNNX. Optimized for mobile and embedded platforms with support for FP16 inference on ARM architectures.
keywords: Ultralytics, NCNN, model export, PyTorch to NCNN, PNNX, mobile deployment, ARM, embedded systems, FP16, lightweight inference, Tencent NCNN, edge AI
---
# Reference for `ultralytics/utils/export/ncnn.py`
!!! success "Improvements"
This page is sourced from [https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/ncnn.py](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/ncnn.py). Have an improvement or example to add? Open a [Pull Request](https://docs.ultralytics.com/help/contributing/) — thank you! 🙏
<br>
## ::: ultralytics.utils.export.ncnn.torch2ncnn
<br><br>

16
docs/en/reference/utils/export/openvino.md Normal file
View file

@ -0,0 +1,16 @@
---
description: OpenVINO export utilities for converting PyTorch YOLO models to OpenVINO format with support for FP16 compression and INT8 quantization via NNCF. Provides optimized inference for Intel hardware including CPUs, GPUs, and VPUs.
keywords: Ultralytics, OpenVINO, model export, PyTorch to OpenVINO, INT8 quantization, FP16, NNCF, Intel, model optimization, inference acceleration, CPU inference, GPU inference, VPU, edge deployment
---
# Reference for `ultralytics/utils/export/openvino.py`
!!! success "Improvements"
This page is sourced from [https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/openvino.py](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/openvino.py). Have an improvement or example to add? Open a [Pull Request](https://docs.ultralytics.com/help/contributing/) — thank you! 🙏
<br>
## ::: ultralytics.utils.export.openvino.torch2openvino
<br><br>

16
docs/en/reference/utils/export/paddle.md Normal file
View file

@ -0,0 +1,16 @@
---
description: PaddlePaddle export utilities for converting PyTorch YOLO models to Paddle format using X2Paddle. Enables deployment on PaddlePaddle inference engine with support for both CPU and GPU backends.
keywords: Ultralytics, PaddlePaddle, model export, PyTorch to Paddle, X2Paddle, Paddle inference, Baidu Paddle, model conversion, CPU inference, GPU inference
---
# Reference for `ultralytics/utils/export/paddle.py`
!!! success "Improvements"
This page is sourced from [https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/paddle.py](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/paddle.py). Have an improvement or example to add? Open a [Pull Request](https://docs.ultralytics.com/help/contributing/) — thank you! 🙏
<br>
## ::: ultralytics.utils.export.paddle.torch2paddle
<br><br>

16
docs/en/reference/utils/export/rknn.md Normal file
View file

@ -0,0 +1,16 @@
---
description: RKNN export utilities for converting ONNX models to RKNN format for Rockchip NPUs. Supports deployment on RK3588 and other Rockchip platforms for optimized inference on edge devices.
keywords: Ultralytics, RKNN, model export, ONNX to RKNN, Rockchip, NPU, RK3588, edge AI, embedded systems, neural network, hardware acceleration
---
# Reference for `ultralytics/utils/export/rknn.py`
!!! success "Improvements"
This page is sourced from [https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/rknn.py](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/rknn.py). Have an improvement or example to add? Open a [Pull Request](https://docs.ultralytics.com/help/contributing/) — thank you! 🙏
<br>
## ::: ultralytics.utils.export.rknn.onnx2rknn
<br><br>

16
docs/en/reference/utils/export/torchscript.md Normal file
View file

@ -0,0 +1,16 @@
---
description: TorchScript export utilities for converting PyTorch YOLO models to TorchScript format. Supports mobile optimization and metadata embedding for production deployment and C++ inference.
keywords: Ultralytics, TorchScript, model export, PyTorch, JIT trace, mobile optimization, production deployment, C++ inference, libtorch, model serialization
---
# Reference for `ultralytics/utils/export/torchscript.py`
!!! success "Improvements"
This page is sourced from [https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/torchscript.py](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/torchscript.py). Have an improvement or example to add? Open a [Pull Request](https://docs.ultralytics.com/help/contributing/) — thank you! 🙏
<br>
## ::: ultralytics.utils.export.torchscript.torch2torchscript
<br><br>

8
mkdocs.yml
View file

@ -828,10 +828,18 @@ nav:
- errors: reference/utils/errors.md
- events: reference/utils/events.md
- export:
- axelera: reference/utils/export/axelera.md
- coreml: reference/utils/export/coreml.md
- engine: reference/utils/export/engine.md
- executorch: reference/utils/export/executorch.md
- imx: reference/utils/export/imx.md
- mnn: reference/utils/export/mnn.md
- ncnn: reference/utils/export/ncnn.md
- openvino: reference/utils/export/openvino.md
- paddle: reference/utils/export/paddle.md
- rknn: reference/utils/export/rknn.md
- tensorflow: reference/utils/export/tensorflow.md
- torchscript: reference/utils/export/torchscript.md
- files: reference/utils/files.md
- git: reference/utils/git.md
- instance: reference/utils/instance.md

581
ultralytics/engine/exporter.py
View file

@ -86,10 +86,8 @@ from ultralytics.nn.tasks import ClassificationModel, DetectionModel, Segmentati
from ultralytics.utils import (
ARM64,
DEFAULT_CFG,
IS_COLAB,
IS_DEBIAN_BOOKWORM,
IS_DEBIAN_TRIXIE,
IS_JETSON,
IS_RASPBERRYPI,
IS_UBUNTU,
LINUX,
@ -119,28 +117,16 @@ from ultralytics.utils.checks import (
is_intel,
is_sudo_available,
)
from ultralytics.utils.export import (
keras2pb,
onnx2engine,
onnx2saved_model,
pb2tfjs,
tflite2edgetpu,
torch2executorch,
torch2imx,
torch2onnx,
)
from ultralytics.utils.files import file_size
from ultralytics.utils.metrics import batch_probiou
from ultralytics.utils.nms import TorchNMS
from ultralytics.utils.ops import Profile
from ultralytics.utils.patches import arange_patch
from ultralytics.utils.torch_utils import (
TORCH_1_10,
TORCH_1_11,
TORCH_1_13,
TORCH_2_1,
TORCH_2_3,
TORCH_2_4,
TORCH_2_9,
select_device,
)
@ -185,36 +171,6 @@ def export_formats():
return dict(zip(["Format", "Argument", "Suffix", "CPU", "GPU", "Arguments"], zip(*x)))
def best_onnx_opset(onnx, cuda=False) -> int:
"""Return max ONNX opset for this torch version with ONNX fallback."""
if TORCH_2_4: # _constants.ONNX_MAX_OPSET first defined in torch 1.13
opset = torch.onnx.utils._constants.ONNX_MAX_OPSET - 1 # use second-latest version for safety
if TORCH_2_9:
opset = min(opset, 20) # legacy TorchScript exporter caps at opset 20 in torch 2.9+
if cuda:
opset -= 2 # fix CUDA ONNXRuntime NMS squeeze op errors
else:
version = ".".join(TORCH_VERSION.split(".")[:2])
opset = {
"1.8": 12,
"1.9": 12,
"1.10": 13,
"1.11": 14,
"1.12": 15,
"1.13": 17,
"2.0": 17, # reduced from 18 to fix ONNX errors
"2.1": 17, # reduced from 19
"2.2": 17, # reduced from 19
"2.3": 17, # reduced from 19
"2.4": 20,
"2.5": 20,
"2.6": 20,
"2.7": 20,
"2.8": 23,
}.get(version, 12)
return min(opset, onnx.defs.onnx_opset_version())
def validate_args(format, passed_args, valid_args):
"""Validate arguments based on the export format.
@ -348,50 +304,27 @@ class Exporter:
raise ValueError(f"{msg} Valid formats are {fmts}")
LOGGER.warning(f"Invalid export format='{fmt}', updating to format='{matches[0]}'")
fmt = matches[0]
flags = [x == fmt for x in fmts]
if sum(flags) != 1:
raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}")
(
jit,
onnx,
xml,
engine,
coreml,
saved_model,
pb,
tflite,
edgetpu,
tfjs,
paddle,
mnn,
ncnn,
imx,
rknn,
executorch,
axelera,
) = flags # export booleans
is_tf_format = any((saved_model, pb, tflite, edgetpu, tfjs))
is_tf_format = fmt in {"saved_model", "pb", "tflite", "edgetpu", "tfjs"}
# Device
dla = None
if engine and self.args.device is None:
self.dla = None
if fmt == "engine" and self.args.device is None:
LOGGER.warning("TensorRT requires GPU export, automatically assigning device=0")
self.args.device = "0"
if engine and "dla" in str(self.args.device): # convert int/list to str first
if fmt == "engine" and "dla" in str(self.args.device): # convert int/list to str first
device_str = str(self.args.device)
dla = device_str.rsplit(":", 1)[-1]
self.dla = device_str.rsplit(":", 1)[-1]
self.args.device = "0" # update device to "0"
assert dla in {"0", "1"}, f"Expected device 'dla:0' or 'dla:1', but got {device_str}."
if imx and self.args.device is None and torch.cuda.is_available():
assert self.dla in {"0", "1"}, f"Expected device 'dla:0' or 'dla:1', but got {device_str}."
if fmt == "imx" and self.args.device is None and torch.cuda.is_available():
LOGGER.warning("Exporting on CPU while CUDA is available, setting device=0 for faster export on GPU.")
self.args.device = "0" # update device to "0"
self.device = select_device("cpu" if self.args.device is None else self.args.device)
# Argument compatibility checks
fmt_keys = fmts_dict["Arguments"][flags.index(True) + 1]
fmt_keys = dict(zip(fmts_dict["Argument"], fmts_dict["Arguments"]))[fmt]
validate_args(fmt, self.args, fmt_keys)
if axelera:
if fmt == "axelera":
if not IS_PYTHON_3_10:
raise SystemError("Axelera export only supported on Python 3.10.")
if not self.args.int8:
@ -401,7 +334,7 @@ class Exporter:
raise ValueError("Axelera export only supported for detection models.")
if not self.args.data:
self.args.data = "coco128.yaml" # Axelera default to coco128.yaml
if imx:
if fmt == "imx":
if not self.args.int8:
LOGGER.warning("IMX export requires int8=True, setting int8=True.")
self.args.int8 = True
@ -418,11 +351,11 @@ class Exporter:
if hasattr(model, "end2end"):
if self.args.end2end is not None:
model.end2end = self.args.end2end
if rknn or ncnn or executorch or paddle or imx or edgetpu:
if fmt in {"rknn", "ncnn", "executorch", "paddle", "imx", "edgetpu"}:
# Disable end2end branch for certain export formats as they does not support topk
model.end2end = False
LOGGER.warning(f"{fmt.upper()} export does not support end2end models, disabling end2end branch.")
if engine and self.args.int8:
if fmt == "engine" and self.args.int8:
# TensorRT<=10.3.0 with int8 has known end2end build issues
# https://github.com/ultralytics/ultralytics/issues/23841
try:
@ -438,16 +371,16 @@ class Exporter:
if self.args.half and self.args.int8:
LOGGER.warning("half=True and int8=True are mutually exclusive, setting half=False.")
self.args.half = False
if self.args.half and jit and self.device.type == "cpu":
if self.args.half and fmt == "torchscript" and self.device.type == "cpu":
LOGGER.warning(
"half=True only compatible with GPU export for TorchScript, i.e. use device=0, setting half=False."
)
self.args.half = False
self.imgsz = check_imgsz(self.args.imgsz, stride=model.stride, min_dim=2) # check image size
if self.args.optimize:
assert not ncnn, "optimize=True not compatible with format='ncnn', i.e. use optimize=False"
assert fmt != "ncnn", "optimize=True not compatible with format='ncnn', i.e. use optimize=False"
assert self.device.type == "cpu", "optimize=True not compatible with cuda devices, i.e. use device='cpu'"
if rknn:
if fmt == "rknn":
if not self.args.name:
LOGGER.warning(
"Rockchip RKNN export requires a missing 'name' arg for processor type. "
@ -460,18 +393,18 @@ class Exporter:
)
if self.args.nms:
assert not isinstance(model, ClassificationModel), "'nms=True' is not valid for classification models."
assert not tflite or not ARM64 or not LINUX, "TFLite export with NMS unsupported on ARM64 Linux"
assert fmt != "tflite" or not ARM64 or not LINUX, "TFLite export with NMS unsupported on ARM64 Linux"
assert not is_tf_format or TORCH_1_13, "TensorFlow exports with NMS require torch>=1.13"
assert not onnx or TORCH_1_13, "ONNX export with NMS requires torch>=1.13"
assert fmt != "onnx" or TORCH_1_13, "ONNX export with NMS requires torch>=1.13"
if getattr(model, "end2end", False) or isinstance(model.model[-1], RTDETRDecoder):
LOGGER.warning("'nms=True' is not available for end2end models. Forcing 'nms=False'.")
self.args.nms = False
self.args.conf = self.args.conf or 0.25 # set conf default value for nms export
if (engine or coreml or self.args.nms) and self.args.dynamic and self.args.batch == 1:
if (fmt in {"engine", "coreml"} or self.args.nms) and self.args.dynamic and self.args.batch == 1:
LOGGER.warning(
f"'dynamic=True' model with '{'nms=True' if self.args.nms else f'format={self.args.format}'}' requires max batch size, i.e. 'batch=16'"
)
if edgetpu:
if fmt == "edgetpu":
if not LINUX or ARM64:
raise SystemError(
"Edge TPU export only supported on non-aarch64 Linux. See https://coral.ai/docs/edgetpu/compiler"
@ -492,7 +425,7 @@ class Exporter:
LOGGER.warning(
f"INT8 export requires a missing 'data' arg for calibration. Using default 'data={self.args.data}'."
)
if tfjs and (ARM64 and LINUX):
if fmt == "tfjs" and ARM64 and LINUX:
raise SystemError("TF.js exports are not currently supported on ARM64 Linux")
# Recommend OpenVINO if export and Intel CPU
if SETTINGS.get("openvino_msg"):
@ -519,15 +452,15 @@ class Exporter:
model.float()
model = model.fuse()
if imx:
if fmt == "imx":
from ultralytics.utils.export.imx import FXModel
model = FXModel(model, self.imgsz)
if tflite or edgetpu:
if fmt in {"tflite", "edgetpu"}:
from ultralytics.utils.export.tensorflow import tf_wrapper
model = tf_wrapper(model)
if executorch:
if fmt == "executorch":
from ultralytics.utils.export.executorch import executorch_wrapper
model = executorch_wrapper(model)
@ -542,7 +475,7 @@ class Exporter:
anchors = sum(int(self.imgsz[0] / s) * int(self.imgsz[1] / s) for s in model.stride.tolist())
m.max_det = min(self.args.max_det, anchors)
m.agnostic_nms = self.args.agnostic_nms
m.xyxy = self.args.nms and not coreml
m.xyxy = self.args.nms and fmt != "coreml"
m.shape = None # reset cached shape for new export input size
if hasattr(model, "pe") and hasattr(m, "fuse"): # for YOLOE models
m.fuse(model.pe.to(self.device))
@ -552,8 +485,8 @@ class Exporter:
y = None
for _ in range(2): # dry runs
y = NMSModel(model, self.args)(im) if self.args.nms and not coreml and not imx else model(im)
if self.args.half and (onnx or jit) and self.device.type != "cpu":
y = NMSModel(model, self.args)(im) if self.args.nms and fmt not in {"coreml", "imx"} else model(im)
if self.args.half and fmt in {"onnx", "torchscript"} and self.device.type != "cpu":
im, model = im.half(), model.half() # to FP16
# Assign
@ -584,8 +517,8 @@ class Exporter:
"channels": model.yaml.get("channels", 3),
"end2end": getattr(model, "end2end", False),
} # model metadata
if dla is not None:
self.metadata["dla"] = dla # make sure `AutoBackend` uses correct dla device if it has one
if self.dla is not None:
self.metadata["dla"] = self.dla # make sure `AutoBackend` uses correct dla device if it has one
if model.task == "pose":
self.metadata["kpt_shape"] = model.model[-1].kpt_shape
if hasattr(model, "kpt_names"):
@ -596,48 +529,24 @@ class Exporter:
f"output shape(s) {self.output_shape} ({file_size(file):.1f} MB)"
)
self.run_callbacks("on_export_start")
# Exports
f = [""] * len(fmts) # exported filenames
if jit: # TorchScript
f[0] = self.export_torchscript()
if engine: # TensorRT required before ONNX
f[1] = self.export_engine(dla=dla)
if onnx: # ONNX
f[2] = self.export_onnx()
if xml: # OpenVINO
f[3] = self.export_openvino()
if coreml: # CoreML
f[4] = self.export_coreml()
if is_tf_format: # TensorFlow formats
self.args.int8 |= edgetpu
f[5], keras_model = self.export_saved_model()
if pb or tfjs: # pb prerequisite to tfjs
f[6] = self.export_pb(keras_model=keras_model)
if tflite:
f[7] = self.export_tflite()
if edgetpu:
f[8] = self.export_edgetpu(tflite_model=Path(f[5]) / f"{self.file.stem}_full_integer_quant.tflite")
if tfjs:
f[9] = self.export_tfjs()
if paddle: # PaddlePaddle
f[10] = self.export_paddle()
if mnn: # MNN
f[11] = self.export_mnn()
if ncnn: # NCNN
f[12] = self.export_ncnn()
if imx:
f[13] = self.export_imx()
if rknn:
f[14] = self.export_rknn()
if executorch:
f[15] = self.export_executorch()
if axelera:
f[16] = self.export_axelera()
# Export
if is_tf_format:
self.args.int8 |= fmt == "edgetpu"
f, keras_model = self.export_saved_model()
if fmt in {"pb", "tfjs"}: # pb prerequisite to tfjs
f = self.export_pb(keras_model=keras_model)
if fmt == "tflite":
f = self.export_tflite()
if fmt == "edgetpu":
f = self.export_edgetpu(tflite_model=Path(f) / f"{self.file.stem}_full_integer_quant.tflite")
if fmt == "tfjs":
f = self.export_tfjs()
else:
f = getattr(self, f"export_{fmt}")()
# Finish
f = [str(x) for x in f if x] # filter out '' and None
if any(f):
f = str(Path(f[-1]))
if f:
square = self.imgsz[0] == self.imgsz[1]
s = (
""
@ -688,19 +597,16 @@ class Exporter:
@try_export
def export_torchscript(self, prefix=colorstr("TorchScript:")):
"""Export YOLO model to TorchScript format."""
LOGGER.info(f"\n{prefix} starting export with torch {TORCH_VERSION}...")
f = self.file.with_suffix(".torchscript")
from ultralytics.utils.export.torchscript import torch2torchscript
ts = torch.jit.trace(NMSModel(self.model, self.args) if self.args.nms else self.model, self.im, strict=False)
extra_files = {"config.txt": json.dumps(self.metadata)} # torch._C.ExtraFilesMap()
if self.args.optimize: # https://pytorch.org/tutorials/recipes/mobile_interpreter.html
LOGGER.info(f"{prefix} optimizing for mobile...")
from torch.utils.mobile_optimizer import optimize_for_mobile
optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files)
else:
ts.save(str(f), _extra_files=extra_files)
return f
return torch2torchscript(
NMSModel(self.model, self.args) if self.args.nms else self.model,
self.im,
self.file,
optimize=self.args.optimize,
metadata=self.metadata,
prefix=prefix,
)
@try_export
def export_onnx(self, prefix=colorstr("ONNX:")):
@ -711,6 +617,8 @@ class Exporter:
check_requirements(requirements)
import onnx
from ultralytics.utils.export.engine import best_onnx_opset, torch2onnx
opset = self.args.opset or best_onnx_opset(onnx, cuda="cuda" in self.device.type)
LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__} opset {opset}...")
if self.args.nms:
@ -732,7 +640,7 @@ class Exporter:
self.args.opset = opset # for NMSModel
self.args.simplify = True # fix OBB runtime error related to topk
with arange_patch(self.args):
with arange_patch(dynamic=bool(dynamic), half=self.args.half, fmt=self.args.format):
torch2onnx(
NMSModel(self.model, self.args) if self.args.nms else self.model,
self.im,
@ -783,17 +691,13 @@ class Exporter:
@try_export
def export_openvino(self, prefix=colorstr("OpenVINO:")):
"""Export YOLO model to OpenVINO format."""
from ultralytics.utils.export import torch2openvino
# OpenVINO <= 2025.1.0 error on macOS 15.4+: https://github.com/openvinotoolkit/openvino/issues/30023"
check_requirements("openvino>=2025.2.0" if MACOS and MACOS_VERSION >= "15.4" else "openvino>=2024.0.0")
import openvino as ov
LOGGER.info(f"\n{prefix} starting export with openvino {ov.__version__}...")
assert TORCH_2_1, f"OpenVINO export requires torch>=2.1 but torch=={TORCH_VERSION} is installed"
ov_model = ov.convert_model(
NMSModel(self.model, self.args) if self.args.nms else self.model,
input=None if self.args.dynamic else [self.im.shape],
example_input=self.im,
)
def serialize(ov_model, file):
"""Set RT info, serialize, and save metadata YAML."""
@ -809,16 +713,13 @@ class Exporter:
ov.save_model(ov_model, file, compress_to_fp16=self.args.half)
YAML.save(Path(file).parent / "metadata.yaml", self.metadata) # add metadata.yaml
calibration_dataset, ignored_scope = None, None
if self.args.int8:
fq = str(self.file).replace(self.file.suffix, f"_int8_openvino_model{os.sep}")
fq_ov = str(Path(fq) / self.file.with_suffix(".xml").name)
# INT8 requires nncf, nncf requires packaging>=23.2 https://github.com/openvinotoolkit/nncf/issues/3463
check_requirements("packaging>=23.2") # must be installed first to build nncf wheel
check_requirements("nncf>=2.14.0,<3.0.0" if not TORCH_2_3 else "nncf>=2.14.0")
import nncf
# Generate calibration data for integer quantization
ignored_scope = None
calibration_dataset = nncf.Dataset(self.get_int8_calibration_dataloader(prefix), self._transform_fn)
if isinstance(self.model.model[-1], Detect):
# Includes all Detect subclasses like Segment, Pose, OBB, WorldDetect, YOLOEDetect
head_module_name = ".".join(list(self.model.named_modules())[-1][0].split(".")[:2])
@ -832,16 +733,19 @@ class Exporter:
types=["Sigmoid"],
)
quantized_ov_model = nncf.quantize(
model=ov_model,
calibration_dataset=nncf.Dataset(self.get_int8_calibration_dataloader(prefix), self._transform_fn),
preset=nncf.QuantizationPreset.MIXED,
ignored_scope=ignored_scope,
)
serialize(quantized_ov_model, fq_ov)
return fq
ov_model = torch2openvino(
model=NMSModel(self.model, self.args) if self.args.nms else self.model,
im=self.im,
dynamic=self.args.dynamic,
half=self.args.half,
int8=self.args.int8,
calibration_dataset=calibration_dataset,
ignored_scope=ignored_scope,
prefix=prefix,
)
f = str(self.file).replace(self.file.suffix, f"_openvino_model{os.sep}")
suffix = f"_{'int8_' if self.args.int8 else ''}openvino_model{os.sep}"
f = str(self.file).replace(self.file.suffix, suffix)
f_ov = str(Path(f) / self.file.with_suffix(".xml").name)
serialize(ov_model, f_ov)
@ -850,97 +754,45 @@ class Exporter:
@try_export
def export_paddle(self, prefix=colorstr("PaddlePaddle:")):
"""Export YOLO model to PaddlePaddle format."""
assert not IS_JETSON, "Jetson Paddle exports not supported yet"
check_requirements(
(
"paddlepaddle-gpu>=3.0.0,<3.3.0" # pin <3.3.0 https://github.com/PaddlePaddle/Paddle/issues/77340
if torch.cuda.is_available()
else "paddlepaddle==3.0.0" # pin 3.0.0 for ARM64
if ARM64
else "paddlepaddle>=3.0.0,<3.3.0", # pin <3.3.0 https://github.com/PaddlePaddle/Paddle/issues/77340
"x2paddle",
)
)
import x2paddle
from x2paddle.convert import pytorch2paddle
from ultralytics.utils.export.paddle import torch2paddle
LOGGER.info(f"\n{prefix} starting export with X2Paddle {x2paddle.__version__}...")
f = str(self.file).replace(self.file.suffix, f"_paddle_model{os.sep}")
pytorch2paddle(module=self.model, save_dir=f, jit_type="trace", input_examples=[self.im]) # export
YAML.save(Path(f) / "metadata.yaml", self.metadata) # add metadata.yaml
return f
return torch2paddle(self.model, self.im, self.file, self.metadata, prefix)
@try_export
def export_mnn(self, prefix=colorstr("MNN:")):
"""Export YOLO model to MNN format using MNN https://github.com/alibaba/MNN."""
assert TORCH_1_10, "MNN export requires torch>=1.10.0 to avoid segmentation faults"
f_onnx = self.export_onnx() # get onnx model first
from ultralytics.utils.export.mnn import onnx2mnn
check_requirements("MNN>=2.9.6")
import MNN
from MNN.tools import mnnconvert
# Setup and checks
LOGGER.info(f"\n{prefix} starting export with MNN {MNN.version()}...")
assert Path(f_onnx).exists(), f"failed to export ONNX file: {f_onnx}"
f = str(self.file.with_suffix(".mnn")) # MNN model file
args = ["", "-f", "ONNX", "--modelFile", f_onnx, "--MNNModel", f, "--bizCode", json.dumps(self.metadata)]
if self.args.int8:
args.extend(("--weightQuantBits", "8"))
if self.args.half:
args.append("--fp16")
mnnconvert.convert(args)
# remove scratch file for model convert optimize
convert_scratch = Path(self.file.parent / ".__convert_external_data.bin")
if convert_scratch.exists():
convert_scratch.unlink()
return f
f_onnx = self.export_onnx()
return onnx2mnn(
f_onnx, self.file, half=self.args.half, int8=self.args.int8, metadata=self.metadata, prefix=prefix
)
@try_export
def export_ncnn(self, prefix=colorstr("NCNN:")):
"""Export YOLO model to NCNN format using PNNX https://github.com/pnnx/pnnx."""
check_requirements("ncnn", cmds="--no-deps") # no deps to avoid installing opencv-python
check_requirements("pnnx")
import ncnn
import pnnx
from ultralytics.utils.export.ncnn import torch2ncnn
LOGGER.info(f"\n{prefix} starting export with NCNN {ncnn.__version__} and PNNX {pnnx.__version__}...")
f = Path(str(self.file).replace(self.file.suffix, f"_ncnn_model{os.sep}"))
ncnn_args = dict(
ncnnparam=(f / "model.ncnn.param").as_posix(),
ncnnbin=(f / "model.ncnn.bin").as_posix(),
ncnnpy=(f / "model_ncnn.py").as_posix(),
return torch2ncnn(
self.model,
self.im,
self.file,
half=self.args.half,
metadata=self.metadata,
device=self.device,
prefix=prefix,
)
pnnx_args = dict(
ptpath=(f / "model.pt").as_posix(),
pnnxparam=(f / "model.pnnx.param").as_posix(),
pnnxbin=(f / "model.pnnx.bin").as_posix(),
pnnxpy=(f / "model_pnnx.py").as_posix(),
pnnxonnx=(f / "model.pnnx.onnx").as_posix(),
)
f.mkdir(exist_ok=True) # make ncnn_model directory
pnnx.export(self.model, inputs=self.im, **ncnn_args, **pnnx_args, fp16=self.args.half, device=self.device.type)
for f_debug in ("debug.bin", "debug.param", "debug2.bin", "debug2.param", *pnnx_args.values()):
Path(f_debug).unlink(missing_ok=True)
YAML.save(f / "metadata.yaml", self.metadata) # add metadata.yaml
return str(f)
@try_export
def export_coreml(self, prefix=colorstr("CoreML:")):
"""Export YOLO model to CoreML format."""
mlmodel = self.args.format.lower() == "mlmodel" # legacy *.mlmodel export format requested
check_requirements(
["coremltools>=9.0", "numpy>=1.14.5,<=2.3.5"]
) # latest numpy 2.4.0rc1 breaks coremltools exports
from ultralytics.utils.export.coreml import IOSDetectModel, pipeline_coreml, torch2coreml
# latest numpy 2.4.0rc1 breaks coremltools exports
check_requirements(["coremltools>=9.0", "numpy>=1.14.5,<=2.3.5"])
import coremltools as ct
LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...")
assert not WINDOWS, "CoreML export is not supported on Windows, please run on macOS or Linux."
assert TORCH_1_11, "CoreML export requires torch>=1.11"
if self.args.batch > 1:
@ -956,18 +808,13 @@ class Exporter:
if f.is_dir():
shutil.rmtree(f)
classifier_config = None
if self.model.task == "classify":
classifier_config = ct.ClassifierConfig(list(self.model.names.values()))
model = self.model
elif self.model.task == "detect":
if self.model.task == "detect":
model = IOSDetectModel(self.model, self.im, mlprogram=not mlmodel) if self.args.nms else self.model
else:
if self.args.nms:
LOGGER.warning(f"{prefix} 'nms=True' is only available for Detect models like 'yolo26n.pt'.")
# TODO CoreML Segment and Pose model pipelining
model = self.model
ts = torch.jit.trace(model.eval(), self.im, strict=False) # TorchScript model
if self.args.dynamic:
input_shape = ct.Shape(
@ -982,37 +829,30 @@ class Exporter:
else:
inputs = [ct.ImageType("image", shape=self.im.shape, scale=1 / 255, bias=[0.0, 0.0, 0.0])]
# Based on apple's documentation it is better to leave out the minimum_deployment target and let that get set
# Internally based on the model conversion and output type.
# Setting minimum_deployment_target >= iOS16 will require setting compute_precision=ct.precision.FLOAT32.
# iOS16 adds in better support for FP16, but none of the CoreML NMS specifications handle FP16 as input.
ct_model = ct.convert(
ts,
ct_model = torch2coreml(
model=model,
inputs=inputs,
classifier_config=classifier_config,
convert_to="neuralnetwork" if mlmodel else "mlprogram",
im=self.im,
classifier_names=list(self.model.names.values()) if self.model.task == "classify" else None,
mlmodel=mlmodel,
half=self.args.half,
int8=self.args.int8,
metadata=self.metadata,
prefix=prefix,
)
bits, mode = (8, "kmeans") if self.args.int8 else (16, "linear") if self.args.half else (32, None)
if bits < 32:
if "kmeans" in mode:
check_requirements("scikit-learn") # scikit-learn package required for k-means quantization
if mlmodel:
ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode)
elif bits == 8: # mlprogram already quantized to FP16
import coremltools.optimize.coreml as cto
op_config = cto.OpPalettizerConfig(mode="kmeans", nbits=bits, weight_threshold=512)
config = cto.OptimizationConfig(global_config=op_config)
ct_model = cto.palettize_weights(ct_model, config=config)
if self.args.nms and self.model.task == "detect":
ct_model = self._pipeline_coreml(ct_model, weights_dir=None if mlmodel else ct_model.weights_dir)
ct_model = pipeline_coreml(
ct_model,
weights_dir=None if mlmodel else ct_model.weights_dir,
metadata=self.metadata,
mlmodel=mlmodel,
iou=self.args.iou,
conf=self.args.conf,
agnostic_nms=self.args.agnostic_nms,
prefix=prefix,
)
m = self.metadata # metadata dict
ct_model.short_description = m.pop("description")
ct_model.author = m.pop("author")
ct_model.license = m.pop("license")
ct_model.version = m.pop("version")
ct_model.user_defined_metadata.update({k: str(v) for k, v in m.items()})
if self.model.task == "classify":
ct_model.user_defined_metadata.update({"com.apple.coreml.model.preview.type": "imageClassifier"})
@ -1028,7 +868,7 @@ class Exporter:
return f
@try_export
def export_engine(self, dla=None, prefix=colorstr("TensorRT:")):
def export_engine(self, prefix=colorstr("TensorRT:")):
"""Export YOLO model to TensorRT format https://developer.nvidia.com/tensorrt."""
assert self.im.device.type != "cpu", "export running on CPU but must be on GPU, i.e. use 'device=0'"
f_onnx = self.export_onnx() # run before TRT import https://github.com/ultralytics/ultralytics/issues/7016
@ -1046,6 +886,8 @@ class Exporter:
check_version(trt.__version__, ">=7.0.0", hard=True)
check_version(trt.__version__, "!=10.1.0", msg="https://github.com/ultralytics/ultralytics/pull/14239")
from ultralytics.utils.export.engine import onnx2engine
# Setup and checks
LOGGER.info(f"\n{prefix} starting export with TensorRT {trt.__version__}...")
assert Path(f_onnx).exists(), f"failed to export ONNX file: {f_onnx}"
@ -1058,7 +900,7 @@ class Exporter:
self.args.int8,
self.args.dynamic,
self.im.shape,
dla=dla,
dla=self.dla,
dataset=self.get_int8_calibration_dataloader(prefix) if self.args.int8 else None,
metadata=self.metadata,
verbose=self.args.verbose,
@ -1099,6 +941,8 @@ class Exporter:
verbose=True,
msg="https://github.com/ultralytics/ultralytics/issues/5161",
)
from ultralytics.utils.export.tensorflow import onnx2saved_model
f = Path(str(self.file).replace(self.file.suffix, "_saved_model"))
if f.is_dir():
shutil.rmtree(f) # delete output folder
@ -1138,6 +982,8 @@ class Exporter:
@try_export
def export_pb(self, keras_model, prefix=colorstr("TensorFlow GraphDef:")):
"""Export YOLO model to TensorFlow GraphDef *.pb format https://github.com/leimao/Frozen-Graph-TensorFlow."""
from ultralytics.utils.export.tensorflow import keras2pb
f = self.file.with_suffix(".pb")
keras2pb(keras_model, f, prefix)
return f
@ -1163,7 +1009,7 @@ class Exporter:
"""Export YOLO model to Axelera format."""
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
try:
from axelera import compiler
from axelera.compiler import CompilerConfig
except ImportError:
check_apt_requirements(
["libllvm14", "libgirepository1.0-dev", "pkg-config", "libcairo2-dev", "build-essential", "cmake"]
@ -1174,9 +1020,9 @@ class Exporter:
cmds="--extra-index-url https://software.axelera.ai/artifactory/axelera-runtime-pypi "
"--extra-index-url https://software.axelera.ai/artifactory/axelera-dev-pypi",
)
from axelera.compiler import CompilerConfig
from axelera import compiler
from axelera.compiler import CompilerConfig
from ultralytics.utils.export.axelera import onnx2axelera
self.args.opset = 17 # hardcode opset for Axelera
onnx_path = self.export_onnx()
@ -1204,25 +1050,14 @@ class Exporter:
output_axm_format=True,
model_name=model_name,
)
qmodel = compiler.quantize(
model=onnx_path,
export_path = onnx2axelera(
onnx_file=onnx_path,
compile_config=config,
metadata=self.metadata,
calibration_dataset=self.get_int8_calibration_dataloader(prefix),
config=config,
transform_fn=self._transform_fn,
prefix=prefix,
)
compiler.compile(model=qmodel, config=config, output_dir=export_path)
axm_name = f"{model_name}.axm"
axm_src = Path(axm_name)
axm_dst = export_path / axm_name
if axm_src.exists():
axm_src.replace(axm_dst)
YAML.save(export_path / "metadata.yaml", self.metadata)
return export_path
@try_export
@ -1230,6 +1065,8 @@ class Exporter:
"""Export YOLO model to ExecuTorch *.pte format."""
assert TORCH_2_9, f"ExecuTorch requires torch>=2.9.0 but torch=={TORCH_VERSION} is installed"
check_executorch_requirements()
from ultralytics.utils.export.executorch import torch2executorch
return torch2executorch(self.model, self.file, self.im, metadata=self.metadata, prefix=prefix)
@try_export
@ -1250,6 +1087,8 @@ class Exporter:
check_apt_requirements(["edgetpu-compiler"])
ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().rsplit(maxsplit=1)[-1]
from ultralytics.utils.export.tensorflow import tflite2edgetpu
LOGGER.info(f"\n{prefix} starting export with Edge TPU compiler {ver}...")
tflite2edgetpu(tflite_file=tflite_model, output_dir=tflite_model.parent, prefix=prefix)
f = str(tflite_model).replace(".tflite", "_edgetpu.tflite") # Edge TPU model
@ -1260,6 +1099,7 @@ class Exporter:
def export_tfjs(self, prefix=colorstr("TensorFlow.js:")):
"""Export YOLO model to TensorFlow.js format."""
check_requirements("tensorflowjs")
from ultralytics.utils.export.tensorflow import pb2tfjs
f = str(self.file).replace(self.file.suffix, "_web_model") # js dir
f_pb = str(self.file.with_suffix(".pb")) # *.pb path
@ -1271,30 +1111,11 @@ class Exporter:
@try_export
def export_rknn(self, prefix=colorstr("RKNN:")):
"""Export YOLO model to RKNN format."""
LOGGER.info(f"\n{prefix} starting export with rknn-toolkit2...")
check_requirements("rknn-toolkit2")
check_requirements("onnx<1.19.0") # fix AttributeError: module 'onnx' has no attribute 'mapping'
if IS_COLAB:
# Prevent 'exit' from closing the notebook https://github.com/airockchip/rknn-toolkit2/issues/259
import builtins
builtins.exit = lambda: None
from rknn.api import RKNN
from ultralytics.utils.export.rknn import onnx2rknn
self.args.opset = min(self.args.opset or 19, 19) # rknn-toolkit expects opset<=19
f = self.export_onnx()
export_path = Path(f"{Path(f).stem}_rknn_model")
export_path.mkdir(exist_ok=True)
rknn = RKNN(verbose=False)
rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform=self.args.name)
rknn.load_onnx(model=f)
rknn.build(do_quantization=False) # TODO: Add quantization support
rknn.export_rknn(str(export_path / f"{Path(f).stem}-{self.args.name}.rknn"))
YAML.save(export_path / "metadata.yaml", self.metadata)
return export_path
f_onnx = self.export_onnx()
return onnx2rknn(f_onnx, name=self.args.name, metadata=self.metadata, prefix=prefix)
@try_export
def export_imx(self, prefix=colorstr("IMX:")):
@ -1317,6 +1138,7 @@ class Exporter:
)
check_requirements("imx500-converter[pt]>=3.17.3")
from ultralytics.utils.export.imx import torch2imx
# Install Java>=17
try:
@ -1350,101 +1172,6 @@ class Exporter:
with zipfile.ZipFile(file, "a", zipfile.ZIP_DEFLATED) as zf:
zf.writestr("metadata.json", json.dumps(self.metadata, indent=2))
def _pipeline_coreml(self, model, weights_dir=None, prefix=colorstr("CoreML Pipeline:")):
"""Create CoreML pipeline with NMS for YOLO detection models."""
import coremltools as ct
LOGGER.info(f"{prefix} starting pipeline with coremltools {ct.__version__}...")
# Output shapes
spec = model.get_spec()
outs = list(iter(spec.description.output))
if self.args.format == "mlmodel": # mlmodel doesn't infer shapes automatically
outs[0].type.multiArrayType.shape[:] = self.output_shape[2], self.output_shape[1] - 4
outs[1].type.multiArrayType.shape[:] = self.output_shape[2], 4
# Checks
names = self.metadata["names"]
nx, ny = spec.description.input[0].type.imageType.width, spec.description.input[0].type.imageType.height
nc = outs[0].type.multiArrayType.shape[-1]
if len(names) != nc: # Hack fix for MLProgram NMS bug https://github.com/ultralytics/ultralytics/issues/22309
names = {**names, **{i: str(i) for i in range(len(names), nc)}}
# Model from spec
model = ct.models.MLModel(spec, weights_dir=weights_dir)
# Create NMS protobuf
nms_spec = ct.proto.Model_pb2.Model()
nms_spec.specificationVersion = spec.specificationVersion
for i in range(len(outs)):
decoder_output = model._spec.description.output[i].SerializeToString()
nms_spec.description.input.add()
nms_spec.description.input[i].ParseFromString(decoder_output)
nms_spec.description.output.add()
nms_spec.description.output[i].ParseFromString(decoder_output)
output_names = ["confidence", "coordinates"]
for i, name in enumerate(output_names):
nms_spec.description.output[i].name = name
for i, out in enumerate(outs):
ma_type = nms_spec.description.output[i].type.multiArrayType
ma_type.shapeRange.sizeRanges.add()
ma_type.shapeRange.sizeRanges[0].lowerBound = 0
ma_type.shapeRange.sizeRanges[0].upperBound = -1
ma_type.shapeRange.sizeRanges.add()
ma_type.shapeRange.sizeRanges[1].lowerBound = out.type.multiArrayType.shape[-1]
ma_type.shapeRange.sizeRanges[1].upperBound = out.type.multiArrayType.shape[-1]
del ma_type.shape[:]
nms = nms_spec.nonMaximumSuppression
nms.confidenceInputFeatureName = outs[0].name # 1x507x80
nms.coordinatesInputFeatureName = outs[1].name # 1x507x4
nms.confidenceOutputFeatureName = output_names[0]
nms.coordinatesOutputFeatureName = output_names[1]
nms.iouThresholdInputFeatureName = "iouThreshold"
nms.confidenceThresholdInputFeatureName = "confidenceThreshold"
nms.iouThreshold = self.args.iou
nms.confidenceThreshold = self.args.conf
nms.pickTop.perClass = not self.args.agnostic_nms
nms.stringClassLabels.vector.extend(names.values())
nms_model = ct.models.MLModel(nms_spec)
# Pipeline models together
pipeline = ct.models.pipeline.Pipeline(
input_features=[
("image", ct.models.datatypes.Array(3, ny, nx)),
("iouThreshold", ct.models.datatypes.Double()),
("confidenceThreshold", ct.models.datatypes.Double()),
],
output_features=output_names,
)
pipeline.add_model(model)
pipeline.add_model(nms_model)
# Correct datatypes
pipeline.spec.description.input[0].ParseFromString(model._spec.description.input[0].SerializeToString())
pipeline.spec.description.output[0].ParseFromString(nms_model._spec.description.output[0].SerializeToString())
pipeline.spec.description.output[1].ParseFromString(nms_model._spec.description.output[1].SerializeToString())
# Update metadata
pipeline.spec.specificationVersion = spec.specificationVersion
pipeline.spec.description.metadata.userDefined.update(
{"IoU threshold": str(nms.iouThreshold), "Confidence threshold": str(nms.confidenceThreshold)}
)
# Save the model
model = ct.models.MLModel(pipeline.spec, weights_dir=weights_dir)
model.input_description["image"] = "Input image"
model.input_description["iouThreshold"] = f"(optional) IoU threshold override (default: {nms.iouThreshold})"
model.input_description["confidenceThreshold"] = (
f"(optional) Confidence threshold override (default: {nms.confidenceThreshold})"
)
model.output_description["confidence"] = 'Boxes × Class confidence (see user-defined metadata "classes")'
model.output_description["coordinates"] = "Boxes × [x, y, width, height] (relative to image size)"
LOGGER.info(f"{prefix} pipeline success")
return model
@staticmethod
def _transform_fn(data_item) -> np.ndarray:
"""The transformation function for Axelera/OpenVINO quantization preprocessing."""
@ -1463,40 +1190,6 @@ class Exporter:
callback(self)
class IOSDetectModel(torch.nn.Module):
"""Wrap an Ultralytics YOLO model for Apple iOS CoreML export."""
def __init__(self, model, im, mlprogram=True):
"""Initialize the IOSDetectModel class with a YOLO model and example image.
Args:
model (torch.nn.Module): The YOLO model to wrap.
im (torch.Tensor): Example input tensor with shape (B, C, H, W).
mlprogram (bool): Whether exporting to MLProgram format.
"""
super().__init__()
_, _, h, w = im.shape # batch, channel, height, width
self.model = model
self.nc = len(model.names) # number of classes
self.mlprogram = mlprogram
if w == h:
self.normalize = 1.0 / w # scalar
else:
self.normalize = torch.tensor(
[1.0 / w, 1.0 / h, 1.0 / w, 1.0 / h], # broadcast (slower, smaller)
device=next(model.parameters()).device,
)
def forward(self, x):
"""Normalize predictions of object detection model with input size-dependent factors."""
xywh, cls = self.model(x)[0].transpose(0, 1).split((4, self.nc), 1)
if self.mlprogram and self.nc % 80 != 0: # NMS bug https://github.com/ultralytics/ultralytics/issues/22309
pad_length = int(((self.nc + 79) // 80) * 80) - self.nc # pad class length to multiple of 80
cls = torch.nn.functional.pad(cls, (0, pad_length, 0, 0), "constant", 0)
return cls, xywh * self.normalize
class NMSModel(torch.nn.Module):
"""Model wrapper with embedded NMS for Detect, Segment, Pose and OBB."""

16
ultralytics/utils/export/__init__.py
View file

@ -1,17 +1,33 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from .axelera import onnx2axelera
from .coreml import torch2coreml
from .engine import onnx2engine, torch2onnx
from .executorch import torch2executorch
from .imx import torch2imx
from .mnn import onnx2mnn
from .ncnn import torch2ncnn
from .openvino import torch2openvino
from .paddle import torch2paddle
from .rknn import onnx2rknn
from .tensorflow import keras2pb, onnx2saved_model, pb2tfjs, tflite2edgetpu
from .torchscript import torch2torchscript
__all__ = [
"keras2pb",
"onnx2axelera",
"onnx2engine",
"onnx2mnn",
"onnx2rknn",
"onnx2saved_model",
"pb2tfjs",
"tflite2edgetpu",
"torch2coreml",
"torch2executorch",
"torch2imx",
"torch2ncnn",
"torch2onnx",
"torch2openvino",
"torch2paddle",
"torch2torchscript",
]

70
ultralytics/utils/export/axelera.py Normal file
View file

@ -0,0 +1,70 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
from pathlib import Path
from typing import Any, Callable
from ultralytics.utils import LOGGER, YAML
def onnx2axelera(
onnx_file: str,
compile_config=None,
metadata: dict | None = None,
calibration_dataset: Any | None = None,
transform_fn: Callable | None = None,
prefix: str = "",
):
"""Export an ONNX model to Axelera format.
Args:
onnx_file (str): Path to the source ONNX file (already exported).
compile_config (axelera.compiler.CompilerConfig): Compiler configuration object. If None, a default
``CompilerConfig`` is created.
metadata (dict | None): Metadata saved as ``metadata.yaml``.
calibration_dataset: Dataloader for INT8 calibration.
transform_fn: Transformation function applied to calibration batches.
prefix (str): Prefix for log messages.
Returns:
(Path): Path to the exported ``_axelera_model`` directory.
"""
from axelera import compiler
LOGGER.info(f"\n{prefix} starting export with axelera...")
model_name = Path(onnx_file).stem
export_path = Path(f"{model_name}_axelera_model")
export_path.mkdir(exist_ok=True)
if compile_config is None:
from axelera.compiler import CompilerConfig
compile_config = CompilerConfig(
tiling_depth=6,
split_buffer_promotion=True,
resources_used=0.25,
aipu_cores_used=1,
multicore_mode="batch",
output_axm_format=True,
model_name=model_name,
)
qmodel = compiler.quantize(
model=onnx_file,
calibration_dataset=calibration_dataset,
config=compile_config,
transform_fn=transform_fn,
)
compiler.compile(model=qmodel, config=compile_config, output_dir=export_path)
axm_name = f"{model_name}.axm"
axm_src = Path(axm_name)
axm_dst = export_path / axm_name
if axm_src.exists():
axm_src.replace(axm_dst)
if metadata:
YAML.save(export_path / "metadata.yaml", metadata)
return export_path

242
ultralytics/utils/export/coreml.py Normal file
View file

@ -0,0 +1,242 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
from pathlib import Path
from typing import Any
import torch
import torch.nn as nn
from ultralytics.utils import LOGGER
class IOSDetectModel(nn.Module):
"""Wrap an Ultralytics YOLO model for Apple iOS CoreML export."""
def __init__(self, model: nn.Module, im: torch.Tensor, mlprogram: bool = True):
"""Initialize the IOSDetectModel class with a YOLO model and example image.
Args:
model (nn.Module): The YOLO model to wrap.
im (torch.Tensor): Example input tensor with shape (B, C, H, W).
mlprogram (bool): Whether exporting to MLProgram format.
"""
super().__init__()
_, _, h, w = im.shape # batch, channel, height, width
self.model = model
self.nc = len(model.names) # number of classes
self.mlprogram = mlprogram
if w == h:
self.normalize = 1.0 / w # scalar
else:
self.normalize = torch.tensor(
[1.0 / w, 1.0 / h, 1.0 / w, 1.0 / h], # broadcast (slower, smaller)
device=next(model.parameters()).device,
)
def forward(self, x: torch.Tensor):
"""Normalize predictions of object detection model with input size-dependent factors."""
xywh, cls = self.model(x)[0].transpose(0, 1).split((4, self.nc), 1)
if self.mlprogram and self.nc % 80 != 0: # NMS bug https://github.com/ultralytics/ultralytics/issues/22309
pad_length = int(((self.nc + 79) // 80) * 80) - self.nc # pad class length to multiple of 80
cls = torch.nn.functional.pad(cls, (0, pad_length, 0, 0), "constant", 0)
return cls, xywh * self.normalize
def pipeline_coreml(
model: Any,
output_shape: tuple,
metadata: dict,
mlmodel: bool = False,
iou: float = 0.45,
conf: float = 0.25,
agnostic_nms: bool = False,
weights_dir: Path | str | None = None,
prefix: str = "",
):
"""Create CoreML pipeline with NMS for YOLO detection models.
Args:
model: CoreML model.
output_shape (tuple): Output shape tuple from the exporter.
metadata (dict): Model metadata.
mlmodel (bool): Whether the model is an MLModel (vs MLProgram).
iou (float): IoU threshold for NMS.
conf (float): Confidence threshold for NMS.
agnostic_nms (bool): Whether to use class-agnostic NMS.
weights_dir (Path | str | None): Weights directory for MLProgram models.
prefix (str): Prefix for log messages.
Returns:
CoreML pipeline model.
"""
import coremltools as ct
LOGGER.info(f"{prefix} starting pipeline with coremltools {ct.__version__}...")
spec = model.get_spec()
outs = list(iter(spec.description.output))
if mlmodel: # mlmodel doesn't infer shapes automatically
outs[0].type.multiArrayType.shape[:] = output_shape[2], output_shape[1] - 4
outs[1].type.multiArrayType.shape[:] = output_shape[2], 4
names = metadata["names"]
nx = spec.description.input[0].type.imageType.width
ny = spec.description.input[0].type.imageType.height
nc = outs[0].type.multiArrayType.shape[-1]
if len(names) != nc: # Hack fix for MLProgram NMS bug https://github.com/ultralytics/ultralytics/issues/22309
names = {**names, **{i: str(i) for i in range(len(names), nc)}}
model = ct.models.MLModel(spec, weights_dir=weights_dir)
# Create NMS protobuf
nms_spec = ct.proto.Model_pb2.Model()
nms_spec.specificationVersion = spec.specificationVersion
for i in range(len(outs)):
decoder_output = model._spec.description.output[i].SerializeToString()
nms_spec.description.input.add()
nms_spec.description.input[i].ParseFromString(decoder_output)
nms_spec.description.output.add()
nms_spec.description.output[i].ParseFromString(decoder_output)
output_names = ["confidence", "coordinates"]
for i, name in enumerate(output_names):
nms_spec.description.output[i].name = name
for i, out in enumerate(outs):
ma_type = nms_spec.description.output[i].type.multiArrayType
ma_type.shapeRange.sizeRanges.add()
ma_type.shapeRange.sizeRanges[0].lowerBound = 0
ma_type.shapeRange.sizeRanges[0].upperBound = -1
ma_type.shapeRange.sizeRanges.add()
ma_type.shapeRange.sizeRanges[1].lowerBound = out.type.multiArrayType.shape[-1]
ma_type.shapeRange.sizeRanges[1].upperBound = out.type.multiArrayType.shape[-1]
del ma_type.shape[:]
nms = nms_spec.nonMaximumSuppression
nms.confidenceInputFeatureName = outs[0].name # 1x507x80
nms.coordinatesInputFeatureName = outs[1].name # 1x507x4
nms.confidenceOutputFeatureName = output_names[0]
nms.coordinatesOutputFeatureName = output_names[1]
nms.iouThresholdInputFeatureName = "iouThreshold"
nms.confidenceThresholdInputFeatureName = "confidenceThreshold"
nms.iouThreshold = iou
nms.confidenceThreshold = conf
nms.pickTop.perClass = not agnostic_nms
nms.stringClassLabels.vector.extend(names.values())
nms_model = ct.models.MLModel(nms_spec)
# Pipeline models together
pipeline = ct.models.pipeline.Pipeline(
input_features=[
("image", ct.models.datatypes.Array(3, ny, nx)),
("iouThreshold", ct.models.datatypes.Double()),
("confidenceThreshold", ct.models.datatypes.Double()),
],
output_features=output_names,
)
pipeline.add_model(model)
pipeline.add_model(nms_model)
# Correct datatypes
pipeline.spec.description.input[0].ParseFromString(model._spec.description.input[0].SerializeToString())
pipeline.spec.description.output[0].ParseFromString(nms_model._spec.description.output[0].SerializeToString())
pipeline.spec.description.output[1].ParseFromString(nms_model._spec.description.output[1].SerializeToString())
# Update metadata
pipeline.spec.specificationVersion = spec.specificationVersion
pipeline.spec.description.metadata.userDefined.update(
{"IoU threshold": str(nms.iouThreshold), "Confidence threshold": str(nms.confidenceThreshold)}
)
# Save the model
model = ct.models.MLModel(pipeline.spec, weights_dir=weights_dir)
model.input_description["image"] = "Input image"
model.input_description["iouThreshold"] = f"(optional) IoU threshold override (default: {nms.iouThreshold})"
model.input_description["confidenceThreshold"] = (
f"(optional) Confidence threshold override (default: {nms.confidenceThreshold})"
)
model.output_description["confidence"] = 'Boxes × Class confidence (see user-defined metadata "classes")'
model.output_description["coordinates"] = "Boxes × [x, y, width, height] (relative to image size)"
LOGGER.info(f"{prefix} pipeline success")
return model
def torch2coreml(
model: nn.Module,
inputs: list,
im: torch.Tensor,
classifier_names: list[str] | None,
coreml_file: Path | str | None = None,
mlmodel: bool = False,
half: bool = False,
int8: bool = False,
metadata: dict | None = None,
prefix: str = "",
):
"""Export a PyTorch model to CoreML ``.mlpackage`` or ``.mlmodel`` format.
Args:
model (nn.Module): The PyTorch model to export.
inputs (list): CoreML input descriptions for the model.
im (torch.Tensor): Example input tensor for tracing.
classifier_names (list[str] | None): Class names for classifier config, or None if not a classifier.
coreml_file (Path | str | None): Output file path, or None to skip saving.
mlmodel (bool): Whether to export as ``.mlmodel`` (neural network) instead of ``.mlpackage`` (ML program).
half (bool): Whether to quantize to FP16.
int8 (bool): Whether to quantize to INT8.
metadata (dict | None): Metadata to embed in the CoreML model.
prefix (str): Prefix for log messages.
Returns:
(ct.models.MLModel): The converted CoreML model.
"""
import coremltools as ct
LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...")
ts = torch.jit.trace(model.eval(), im, strict=False) # TorchScript model
# Based on apple's documentation it is better to leave out the minimum_deployment target and let that get set
# Internally based on the model conversion and output type.
# Setting minimum_deployment_target >= iOS16 will require setting compute_precision=ct.precision.FLOAT32.
# iOS16 adds in better support for FP16, but none of the CoreML NMS specifications handle FP16 as input.
ct_model = ct.convert(
ts,
inputs=inputs,
classifier_config=ct.ClassifierConfig(classifier_names) if classifier_names else None,
convert_to="neuralnetwork" if mlmodel else "mlprogram",
)
bits, mode = (8, "kmeans") if int8 else (16, "linear") if half else (32, None)
if bits < 32:
if "kmeans" in mode:
from ultralytics.utils.checks import check_requirements
check_requirements("scikit-learn") # scikit-learn package required for k-means quantization
if mlmodel:
ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode)
elif bits == 8: # mlprogram already quantized to FP16
import coremltools.optimize.coreml as cto
op_config = cto.OpPalettizerConfig(mode="kmeans", nbits=bits, weight_threshold=512)
config = cto.OptimizationConfig(global_config=op_config)
ct_model = cto.palettize_weights(ct_model, config=config)
m = dict(metadata or {}) # copy to avoid mutating original
ct_model.short_description = m.pop("description", "")
ct_model.author = m.pop("author", "")
ct_model.license = m.pop("license", "")
ct_model.version = m.pop("version", "")
ct_model.user_defined_metadata.update({k: str(v) for k, v in m.items()})
if coreml_file is not None:
try:
ct_model.save(str(coreml_file)) # save *.mlpackage
except Exception as e:
LOGGER.warning(
f"{prefix} CoreML export to *.mlpackage failed ({e}), reverting to *.mlmodel export. "
f"Known coremltools Python 3.11 and Windows bugs https://github.com/apple/coremltools/issues/1928."
)
coreml_file = Path(coreml_file).with_suffix(".mlmodel")
ct_model.save(str(coreml_file))
return ct_model

35
ultralytics/utils/export/engine.py
View file

@ -3,12 +3,43 @@
from __future__ import annotations
import json
import types
from pathlib import Path
import torch
from ultralytics.utils import IS_JETSON, LOGGER
from ultralytics.utils.torch_utils import TORCH_2_4
from ultralytics.utils import IS_JETSON, LOGGER, TORCH_VERSION
from ultralytics.utils.torch_utils import TORCH_2_4, TORCH_2_9
def best_onnx_opset(onnx: types.ModuleType, cuda: bool = False) -> int:
"""Return max ONNX opset for this torch version with ONNX fallback."""
if TORCH_2_4: # _constants.ONNX_MAX_OPSET first defined in torch 1.13
opset = torch.onnx.utils._constants.ONNX_MAX_OPSET - 1 # use second-latest version for safety
if TORCH_2_9:
opset = min(opset, 20) # legacy TorchScript exporter caps at opset 20 in torch 2.9+
if cuda:
opset -= 2 # fix CUDA ONNXRuntime NMS squeeze op errors
else:
version = ".".join(TORCH_VERSION.split(".")[:2])
opset = {
"1.8": 12,
"1.9": 12,
"1.10": 13,
"1.11": 14,
"1.12": 15,
"1.13": 17,
"2.0": 17, # reduced from 18 to fix ONNX errors
"2.1": 17, # reduced from 19
"2.2": 17, # reduced from 19
"2.3": 17, # reduced from 19
"2.4": 20,
"2.5": 20,
"2.6": 20,
"2.7": 20,
"2.8": 23,
}.get(version, 12)
return min(opset, onnx.defs.onnx_opset_version())
def torch2onnx(

55
ultralytics/utils/export/mnn.py Normal file
View file

@ -0,0 +1,55 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
import json
from pathlib import Path
from ultralytics.utils import LOGGER
def onnx2mnn(
f_onnx: str,
file: Path | str,
half: bool = False,
int8: bool = False,
metadata: dict | None = None,
prefix: str = "",
) -> str:
"""Convert an ONNX model to MNN format.
Args:
f_onnx (str): Path to the source ONNX file.
file (Path | str): Source model path used to derive the output ``.mnn`` path.
half (bool): Whether to enable FP16 conversion.
int8 (bool): Whether to enable INT8 weight quantization.
metadata (dict | None): Optional metadata embedded via ``--bizCode``.
prefix (str): Prefix for log messages.
Returns:
(str): Path to the exported ``.mnn`` file.
"""
from ultralytics.utils.checks import check_requirements
from ultralytics.utils.torch_utils import TORCH_1_10
assert TORCH_1_10, "MNN export requires torch>=1.10.0 to avoid segmentation faults"
assert Path(f_onnx).exists(), f"failed to export ONNX file: {f_onnx}"
check_requirements("MNN>=2.9.6")
import MNN
from MNN.tools import mnnconvert
LOGGER.info(f"\n{prefix} starting export with MNN {MNN.version()}...")
file = Path(file)
f = str(file.with_suffix(".mnn")) # MNN model file
mnn_args = ["", "-f", "ONNX", "--modelFile", f_onnx, "--MNNModel", f, "--bizCode", json.dumps(metadata or {})]
if int8:
mnn_args.extend(("--weightQuantBits", "8"))
if half:
mnn_args.append("--fp16")
mnnconvert.convert(mnn_args)
# Remove scratch file created during model convert optimize
convert_scratch = file.parent / ".__convert_external_data.bin"
if convert_scratch.exists():
convert_scratch.unlink()
return f

69
ultralytics/utils/export/ncnn.py Normal file
View file

@ -0,0 +1,69 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
import os
from pathlib import Path
import torch
from ultralytics.utils import LOGGER, YAML
def torch2ncnn(
model: torch.nn.Module,
im: torch.Tensor,
file: Path | str,
half: bool = False,
metadata: dict | None = None,
device: torch.device | None = None,
prefix: str = "",
) -> str:
"""Export a PyTorch model to NCNN format using PNNX.
Args:
model (torch.nn.Module): The PyTorch model to export.
im (torch.Tensor): Example input tensor for tracing.
file (Path | str): Source model path used to derive the output directory.
half (bool): Whether to enable FP16 export.
metadata (dict | None): Optional metadata saved as ``metadata.yaml``.
device (torch.device | None): Device the model lives on.
prefix (str): Prefix for log messages.
Returns:
(str): Path to the exported ``_ncnn_model`` directory.
"""
from ultralytics.utils.checks import check_requirements
check_requirements("ncnn", cmds="--no-deps") # no deps to avoid installing opencv-python
check_requirements("pnnx")
import ncnn
import pnnx
LOGGER.info(f"\n{prefix} starting export with NCNN {ncnn.__version__} and PNNX {pnnx.__version__}...")
file = Path(file)
f = Path(str(file).replace(file.suffix, f"_ncnn_model{os.sep}"))
ncnn_args = dict(
ncnnparam=(f / "model.ncnn.param").as_posix(),
ncnnbin=(f / "model.ncnn.bin").as_posix(),
ncnnpy=(f / "model_ncnn.py").as_posix(),
)
pnnx_args = dict(
ptpath=(f / "model.pt").as_posix(),
pnnxparam=(f / "model.pnnx.param").as_posix(),
pnnxbin=(f / "model.pnnx.bin").as_posix(),
pnnxpy=(f / "model_pnnx.py").as_posix(),
pnnxonnx=(f / "model.pnnx.onnx").as_posix(),
)
f.mkdir(exist_ok=True) # make ncnn_model directory
device_type = device.type if device is not None else "cpu"
pnnx.export(model, inputs=im, **ncnn_args, **pnnx_args, fp16=half, device=device_type)
for f_debug in ("debug.bin", "debug.param", "debug2.bin", "debug2.param", *pnnx_args.values()):
Path(f_debug).unlink(missing_ok=True)
if metadata:
YAML.save(f / "metadata.yaml", metadata) # add metadata.yaml
return str(f)

62
ultralytics/utils/export/openvino.py Normal file
View file

@ -0,0 +1,62 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
import os
from pathlib import Path
from typing import Any
import torch
from ultralytics.utils import LOGGER
def torch2openvino(
model: torch.nn.Module,
im: torch.Tensor,
file: Path | str | None = None,
dynamic: bool = False,
half: bool = False,
int8: bool = False,
calibration_dataset: Any | None = None,
ignored_scope: dict | None = None,
prefix: str = "",
) -> str:
"""Export a PyTorch model to OpenVINO format with optional INT8 quantization.
Args:
model (torch.nn.Module): The model to export (may be NMS-wrapped).
im (torch.Tensor): Example input tensor.
file (Path | str | None): Source model path used to derive output directory.
dynamic (bool): Whether to use dynamic input shapes.
half (bool): Whether to compress to FP16.
int8 (bool): Whether to apply INT8 quantization.
calibration_dataset (nn.Dataset): Dataset for nncf.Dataset (required when ``int8=True``).
ignored_scope (dict | None): Kwargs passed to ``nncf.IgnoredScope`` for head patterns.
prefix (str): Prefix for log messages.
Returns:
(ov.Model): The converted OpenVINO model.
"""
import openvino as ov
LOGGER.info(f"\n{prefix} starting export with openvino {ov.__version__}...")
ov_model = ov.convert_model(model, input=None if dynamic else [im.shape], example_input=im)
if int8:
import nncf
ov_model = nncf.quantize(
model=ov_model,
calibration_dataset=calibration_dataset,
preset=nncf.QuantizationPreset.MIXED,
ignored_scope=ignored_scope,
)
if file is not None:
file = Path(file)
suffix = f"_{'int8_' if int8 else ''}openvino_model{os.sep}"
f = str(file).replace(file.suffix, suffix)
f_ov = str(Path(f) / file.with_suffix(".xml").name)
ov.save_model(ov_model, f_ov, compress_to_fp16=half)
return ov_model

56
ultralytics/utils/export/paddle.py Normal file
View file

@ -0,0 +1,56 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
import os
from pathlib import Path
import torch
from ultralytics.utils import ARM64, IS_JETSON, LOGGER, YAML
def torch2paddle(
model: torch.nn.Module,
im: torch.Tensor,
file: Path | str,
metadata: dict | None = None,
prefix: str = "",
) -> str:
"""Export a PyTorch model to PaddlePaddle format using X2Paddle.
Args:
model (torch.nn.Module): The PyTorch model to export.
im (torch.Tensor): Example input tensor for tracing.
file (Path | str): Source model path used to derive the output directory.
metadata (dict | None): Optional metadata saved as ``metadata.yaml``.
prefix (str): Prefix for log messages.
Returns:
(str): Path to the exported ``_paddle_model`` directory.
"""
assert not IS_JETSON, "Jetson Paddle exports not supported yet"
from ultralytics.utils.checks import check_requirements
check_requirements(
(
"paddlepaddle-gpu>=3.0.0,<3.3.0" # pin <3.3.0 https://github.com/PaddlePaddle/Paddle/issues/77340
if torch.cuda.is_available()
else "paddlepaddle==3.0.0" # pin 3.0.0 for ARM64
if ARM64
else "paddlepaddle>=3.0.0,<3.3.0", # pin <3.3.0 https://github.com/PaddlePaddle/Paddle/issues/77340
"x2paddle",
)
)
import x2paddle
from x2paddle.convert import pytorch2paddle
LOGGER.info(f"\n{prefix} starting export with X2Paddle {x2paddle.__version__}...")
file = Path(file)
f = str(file).replace(file.suffix, f"_paddle_model{os.sep}")
pytorch2paddle(module=model, save_dir=f, jit_type="trace", input_examples=[im]) # export
if metadata:
YAML.save(Path(f) / "metadata.yaml", metadata) # add metadata.yaml
return f

51
ultralytics/utils/export/rknn.py Normal file
View file

@ -0,0 +1,51 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
from pathlib import Path
from ultralytics.utils import IS_COLAB, LOGGER, YAML
def onnx2rknn(
f_onnx: str,
name: str = "rk3588",
metadata: dict | None = None,
prefix: str = "",
) -> Path:
"""Export an ONNX model to RKNN format for Rockchip NPUs.
Args:
f_onnx (str): Path to the source ONNX file (already exported, opset <=19).
name (str): Target platform name (e.g. ``"rk3588"``).
metadata (dict | None): Metadata saved as ``metadata.yaml``.
prefix (str): Prefix for log messages.
Returns:
(Path): Path to the exported ``_rknn_model`` directory.
"""
from ultralytics.utils.checks import check_requirements
LOGGER.info(f"\n{prefix} starting export with rknn-toolkit2...")
check_requirements("rknn-toolkit2")
check_requirements("onnx<1.19.0") # fix AttributeError: module 'onnx' has no attribute 'mapping'
if IS_COLAB:
# Prevent 'exit' from closing the notebook https://github.com/airockchip/rknn-toolkit2/issues/259
import builtins
builtins.exit = lambda: None
from rknn.api import RKNN
export_path = Path(f"{Path(f_onnx).stem}_rknn_model")
export_path.mkdir(exist_ok=True)
rknn = RKNN(verbose=False)
rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform=name)
rknn.load_onnx(model=f_onnx)
rknn.build(do_quantization=False) # TODO: Add quantization support
rknn.export_rknn(str(export_path / f"{Path(f_onnx).stem}-{name}.rknn"))
if metadata:
YAML.save(export_path / "metadata.yaml", metadata)
return export_path

47
ultralytics/utils/export/torchscript.py Normal file
View file

@ -0,0 +1,47 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
import json
from pathlib import Path
import torch
from ultralytics.utils import LOGGER, TORCH_VERSION
def torch2torchscript(
model: torch.nn.Module,
im: torch.Tensor,
file: Path | str,
optimize: bool = False,
metadata: dict | None = None,
prefix: str = "",
) -> Path:
"""Export a PyTorch model to TorchScript format.
Args:
model (torch.nn.Module): The PyTorch model to export (may be NMS-wrapped).
im (torch.Tensor): Example input tensor for tracing.
file (Path | str): Source model file path used to derive output path.
optimize (bool): Whether to optimize for mobile deployment.
metadata (dict | None): Optional metadata to embed in the TorchScript archive.
prefix (str): Prefix for log messages.
Returns:
(Path): Path to the exported ``.torchscript`` file.
"""
LOGGER.info(f"\n{prefix} starting export with torch {TORCH_VERSION}...")
file = Path(file)
f = file.with_suffix(".torchscript")
ts = torch.jit.trace(model, im, strict=False)
extra_files = {"config.txt": json.dumps(metadata or {})} # torch._C.ExtraFilesMap()
if optimize: # https://pytorch.org/tutorials/recipes/mobile_interpreter.html
LOGGER.info(f"{prefix} optimizing for mobile...")
from torch.utils.mobile_optimizer import optimize_for_mobile
optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files)
else:
ts.save(str(f), _extra_files=extra_files)
return f

4
ultralytics/utils/patches.py
View file

@ -202,12 +202,12 @@ def torch_save(*args, **kwargs):
@contextmanager
def arange_patch(args):
def arange_patch(dynamic: bool = False, half: bool = False, fmt: str = ""):
"""Workaround for ONNX torch.arange incompatibility with FP16.
https://github.com/pytorch/pytorch/issues/148041.
"""
if args.dynamic and args.half and args.format == "onnx":
if dynamic and half and fmt == "onnx":
func = torch.arange
def arange(*args, dtype=None, **kwargs):