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

ultralytics 8.4.38 Unify args naming for standalone export functions (#24120)

Browse source 
Signed-off-by: Jing Qiu <61612323+Laughing-q@users.noreply.github.com>
Signed-off-by: Onuralp SEZER <onuralp@ultralytics.com>
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: Ultralytics Assistant <135830346+UltralyticsAssistant@users.noreply.github.com>
Co-authored-by: Onuralp SEZER <onuralp@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: Lakshantha Dissanayake <lakshantha@ultralytics.com>
This commit is contained in:
Jing Qiu 2026-04-16 20:00:28 +08:00 committed by GitHub
parent dd7f30e51e
commit 3108aa614d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 251 additions and 196 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
docs/en/integrations/sony-imx500.md
View file

@ -205,9 +205,9 @@ The export process will create an ONNX model for quantization validation, along
├── dnnParams.xml
├── labels.txt
├── packerOut.zip
├── yolo11n_imx.onnx
├── yolo11n_imx_MemoryReport.json
└── yolo11n_imx.pbtxt
├── model_imx.onnx
├── model_imx_MemoryReport.json
└── model_imx.pbtxt
```
=== "Pose Estimation"
@ -217,9 +217,9 @@ The export process will create an ONNX model for quantization validation, along
├── dnnParams.xml
├── labels.txt
├── packerOut.zip
├── yolo11n-pose_imx.onnx
├── yolo11n-pose_imx_MemoryReport.json
└── yolo11n-pose_imx.pbtxt
├── model_imx.onnx
├── model_imx_MemoryReport.json
└── model_imx.pbtxt
```
=== "Classification"
@ -229,9 +229,9 @@ The export process will create an ONNX model for quantization validation, along
├── dnnParams.xml
├── labels.txt
├── packerOut.zip
├── yolo11n-cls_imx.onnx
├── yolo11n-cls_imx_MemoryReport.json
└── yolo11n-cls_imx.pbtxt
├── model_imx.onnx
├── model_imx_MemoryReport.json
└── model_imx.pbtxt
```
=== "Instance Segmentation"
@ -241,9 +241,9 @@ The export process will create an ONNX model for quantization validation, along
├── dnnParams.xml
├── labels.txt
├── packerOut.zip
├── yolo11n-seg_imx.onnx
├── yolo11n-seg_imx_MemoryReport.json
└── yolo11n-seg_imx.pbtxt
├── model_imx.onnx
├── model_imx_MemoryReport.json
└── model_imx.pbtxt
```
## Using IMX500 Export in Deployment

5
tests/test_exports.py
View file

@ -341,7 +341,7 @@ def test_export_executorch():
file = YOLO(MODEL).export(format="executorch", imgsz=32)
assert Path(file).exists(), f"ExecuTorch export failed, directory not found: {file}"
# Check that .pte file exists in the exported directory
pte_file = Path(file) / Path(MODEL).with_suffix(".pte").name
pte_file = Path(file) / "model.pte"
assert pte_file.exists(), f"ExecuTorch .pte file not found: {pte_file}"
# Check that metadata.yaml exists
metadata_file = Path(file) / "metadata.yaml"
@ -359,8 +359,7 @@ def test_export_executorch_matrix(task):
file = YOLO(TASK2MODEL[task]).export(format="executorch", imgsz=32)
assert Path(file).exists(), f"ExecuTorch export failed for task '{task}', directory not found: {file}"
# Check that .pte file exists in the exported directory
model_name = Path(TASK2MODEL[task]).with_suffix(".pte").name
pte_file = Path(file) / model_name
pte_file = Path(file) / "model.pte"
assert pte_file.exists(), f"ExecuTorch .pte file not found for task '{task}': {pte_file}"
# Check that metadata.yaml exists
metadata_file = Path(file) / "metadata.yaml"

2
ultralytics/__init__.py
View file

@ -1,6 +1,6 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
__version__ = "8.4.37"
__version__ = "8.4.38"
import importlib
import os

88
ultralytics/engine/exporter.py
View file

@ -601,9 +601,9 @@ class Exporter:
from ultralytics.utils.export.torchscript import torch2torchscript
return torch2torchscript(
NMSModel(self.model, self.args) if self.args.nms else self.model,
self.im,
self.file,
model=NMSModel(self.model, self.args) if self.args.nms else self.model,
im=self.im,
output_file=self.file.with_suffix(".torchscript"),
optimize=self.args.optimize,
metadata=self.metadata,
prefix=prefix,
@ -692,9 +692,9 @@ class Exporter:
@try_export
def export_openvino(self, prefix=colorstr("OpenVINO:")):
"""Export YOLO model to OpenVINO format."""
from ultralytics.utils.export import torch2openvino
from ultralytics.utils.export.openvino import torch2openvino
# OpenVINO <= 2025.1.0 error on macOS 15.4+: https://github.com/openvinotoolkit/openvino/issues/30023"
# 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
@ -757,16 +757,26 @@ class Exporter:
"""Export YOLO model to PaddlePaddle format."""
from ultralytics.utils.export.paddle import torch2paddle
return torch2paddle(self.model, self.im, self.file, self.metadata, prefix)
return torch2paddle(
model=self.model,
im=self.im,
output_dir=str(self.file).replace(self.file.suffix, f"_paddle_model{os.sep}"),
metadata=self.metadata,
prefix=prefix,
)
@try_export
def export_mnn(self, prefix=colorstr("MNN:")):
"""Export YOLO model to MNN format using MNN https://github.com/alibaba/MNN."""
from ultralytics.utils.export.mnn import onnx2mnn
f_onnx = self.export_onnx()
return onnx2mnn(
f_onnx, self.file, half=self.args.half, int8=self.args.int8, metadata=self.metadata, prefix=prefix
onnx_file=self.export_onnx(),
output_file=self.file.with_suffix(".mnn"),
half=self.args.half,
int8=self.args.int8,
metadata=self.metadata,
prefix=prefix,
)
@try_export
@ -775,9 +785,9 @@ class Exporter:
from ultralytics.utils.export.ncnn import torch2ncnn
return torch2ncnn(
self.model,
self.im,
self.file,
model=self.model,
im=self.im,
output_dir=str(self.file).replace(self.file.suffix, "_ncnn_model/"),
half=self.args.half,
metadata=self.metadata,
device=self.device,
@ -986,9 +996,7 @@ class Exporter:
"""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
return keras2pb(keras_model, output_file=self.file.with_suffix(".pb"), prefix=prefix)
@try_export
def export_tflite(self, prefix=colorstr("TensorFlow Lite:")):
@ -1016,11 +1024,13 @@ class Exporter:
from ultralytics.utils.export.axelera import torch2axelera
output_dir = self.file.parent / f"{self.file.stem}_axelera_model"
return torch2axelera(
model=self.model,
file=self.file,
output_dir=output_dir,
calibration_dataset=self.get_int8_calibration_dataloader(prefix),
transform_fn=self._transform_fn,
model_name=self.file.stem,
metadata=self.metadata,
prefix=prefix,
)
@ -1032,7 +1042,13 @@ class Exporter:
check_executorch_requirements()
from ultralytics.utils.export.executorch import torch2executorch
return torch2executorch(self.model, self.file, self.im, metadata=self.metadata, prefix=prefix)
return torch2executorch(
model=self.model,
im=self.im,
output_dir=str(self.file).replace(self.file.suffix, "_executorch_model/"),
metadata=self.metadata,
prefix=prefix,
)
@try_export
def export_edgetpu(self, tflite_model="", prefix=colorstr("Edge TPU:")):
@ -1055,10 +1071,9 @@ class Exporter:
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
self._add_tflite_metadata(f)
return f
output_file = tflite2edgetpu(tflite_file=tflite_model, output_dir=tflite_model.parent, prefix=prefix)
self._add_tflite_metadata(output_file)
return output_file
@try_export
def export_tfjs(self, prefix=colorstr("TensorFlow.js:")):
@ -1066,12 +1081,15 @@ class Exporter:
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
pb2tfjs(pb_file=f_pb, output_dir=f, half=self.args.half, int8=self.args.int8, prefix=prefix)
# Add metadata
YAML.save(Path(f) / "metadata.yaml", self.metadata) # add metadata.yaml
return f
output_dir = pb2tfjs(
pb_file=str(self.file.with_suffix(".pb")),
output_dir=str(self.file).replace(self.file.suffix, "_web_model/"),
half=self.args.half,
int8=self.args.int8,
prefix=prefix,
)
YAML.save(Path(output_dir) / "metadata.yaml", self.metadata)
return output_dir
@try_export
def export_rknn(self, prefix=colorstr("RKNN:")):
@ -1080,7 +1098,13 @@ class Exporter:
self.args.opset = min(self.args.opset or 19, 19) # rknn-toolkit expects opset<=19
f_onnx = self.export_onnx()
return onnx2rknn(f_onnx, name=self.args.name, metadata=self.metadata, prefix=prefix)
return onnx2rknn(
onnx_file=f_onnx,
output_dir=str(self.file).replace(self.file.suffix, f"_rknn_model{os.sep}"),
name=self.args.name,
metadata=self.metadata,
prefix=prefix,
)
@try_export
def export_imx(self, prefix=colorstr("IMX:")):
@ -1120,11 +1144,11 @@ class Exporter:
check_apt_requirements(["openjdk-17-jre"])
return torch2imx(
self.model,
self.file,
self.args.conf,
self.args.iou,
self.args.max_det,
model=self.model,
output_dir=str(self.file).replace(self.file.suffix, "_imx_model/"),
conf=self.args.conf,
iou=self.args.iou,
max_det=self.args.max_det,
metadata=self.metadata,
dataset=self.get_int8_calibration_dataloader(prefix),
prefix=prefix,

6
ultralytics/nn/backends/coreml.py
View file

@ -32,7 +32,9 @@ class CoreMLBackend(BaseBackend):
LOGGER.info(f"Loading {weight} for CoreML inference...")
self.model = ct.models.MLModel(weight)
self.dynamic = self.model.get_spec().description.input[0].type.HasField("multiArrayType")
spec = self.model.get_spec()
self.input_name = spec.description.input[0].name
self.dynamic = spec.description.input[0].type.HasField("multiArrayType")
# Load metadata
self.apply_metadata(dict(self.model.user_defined_metadata))
@ -50,7 +52,7 @@ class CoreMLBackend(BaseBackend):
h, w = im.shape[1:3]
im = im.transpose(0, 3, 1, 2) if self.dynamic else Image.fromarray((im[0] * 255).astype("uint8"))
y = self.model.predict({"image": im})
y = self.model.predict({self.input_name: im})
if "confidence" in y: # NMS included
from ultralytics.utils.ops import xywh2xyxy

26
ultralytics/utils/export/axelera.py
View file

@ -16,24 +16,26 @@ from ultralytics.utils.checks import check_requirements
def torch2axelera(
model: torch.nn.Module,
file: str | Path,
output_dir: Path | str,
calibration_dataset: torch.utils.data.DataLoader,
transform_fn: Callable[[Any], np.ndarray],
model_name: str = "model",
metadata: dict | None = None,
prefix: str = "",
) -> Path:
) -> str:
"""Convert a YOLO model to Axelera format.
Args:
model (torch.nn.Module): Source YOLO model for quantization.
file (str | Path): Source model file path used to derive output names.
output_dir (Path | str): Directory to save the exported Axelera model.
calibration_dataset (torch.utils.data.DataLoader): Calibration dataloader for quantization.
transform_fn (Callable[[Any], np.ndarray]): Calibration preprocessing transform function.
model_name (str, optional): Name for the compiled model. Defaults to "model".
metadata (dict | None, optional): Optional metadata to save as YAML. Defaults to None.
prefix (str, optional): Prefix for log messages. Defaults to "".
Returns:
(Path): Path to exported Axelera model directory.
(str): Path to exported Axelera model directory.
"""
prev_protobuf = os.environ.get("PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION")
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
@ -51,10 +53,8 @@ def torch2axelera(
LOGGER.info(f"\n{prefix} starting export with Axelera compiler...")
file = Path(file)
model_name = file.stem
export_path = Path(f"{model_name}_axelera_model")
export_path.mkdir(exist_ok=True)
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
axelera_model_metadata = extract_ultralytics_metadata(model)
config = CompilerConfig(
@ -71,22 +71,22 @@ def torch2axelera(
config=config,
transform_fn=transform_fn,
)
compiler.compile(model=qmodel, config=config, output_dir=export_path)
compiler.compile(model=qmodel, config=config, output_dir=output_dir)
for artifact in [f"{model_name}.axm", "compiler_config_final.toml"]:
artifact_path = Path(artifact)
if artifact_path.exists():
artifact_path.replace(export_path / artifact_path.name)
artifact_path.replace(output_dir / artifact_path.name)
# Remove intermediate compiler artifacts, keeping only the compiled model and config.
keep_suffixes = {".axm"}
keep_names = {"compiler_config_final.toml", "metadata.yaml"}
for f in export_path.iterdir():
for f in output_dir.iterdir():
if f.is_file() and f.suffix not in keep_suffixes and f.name not in keep_names:
f.unlink()
if metadata is not None:
YAML.save(export_path / "metadata.yaml", metadata)
YAML.save(output_dir / "metadata.yaml", metadata)
# Restore original PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION value
if prev_protobuf is None:
@ -94,4 +94,4 @@ def torch2axelera(
else:
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = prev_protobuf
return export_path
return str(output_dir)

18
ultralytics/utils/export/coreml.py
View file

@ -46,7 +46,7 @@ class IOSDetectModel(nn.Module):
def pipeline_coreml(
model: Any,
output_shape: tuple,
output_shape: tuple[int, ...],
metadata: dict,
mlmodel: bool = False,
iou: float = 0.45,
@ -59,7 +59,7 @@ def pipeline_coreml(
Args:
model: CoreML model.
output_shape (tuple): Output shape tuple from the exporter.
output_shape (tuple[int, ...]): 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.
@ -168,13 +168,13 @@ def torch2coreml(
inputs: list,
im: torch.Tensor,
classifier_names: list[str] | None,
coreml_file: Path | str | None = None,
output_file: Path | str | None = None,
mlmodel: bool = False,
half: bool = False,
int8: bool = False,
metadata: dict | None = None,
prefix: str = "",
):
) -> Any:
"""Export a PyTorch model to CoreML ``.mlpackage`` or ``.mlmodel`` format.
Args:
@ -182,7 +182,7 @@ def torch2coreml(
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.
output_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.
@ -229,14 +229,14 @@ def torch2coreml(
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:
if output_file is not None:
try:
ct_model.save(str(coreml_file)) # save *.mlpackage
ct_model.save(str(output_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))
output_file = Path(output_file).with_suffix(".mlmodel")
ct_model.save(str(output_file))
return ct_model

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

@ -44,46 +44,54 @@ def best_onnx_opset(onnx: types.ModuleType, cuda: bool = False) -> int:
@ThreadingLocked()
def torch2onnx(
torch_model: torch.nn.Module,
im: torch.Tensor,
onnx_file: str,
model: torch.nn.Module,
im: torch.Tensor | tuple[torch.Tensor, ...],
output_file: Path | str,
opset: int = 14,
input_names: list[str] = ["images"],
output_names: list[str] = ["output0"],
dynamic: bool | dict = False,
) -> None:
input_names: list[str] | None = None,
output_names: list[str] | None = None,
dynamic: dict | None = None,
) -> str:
"""Export a PyTorch model to ONNX format.
Args:
torch_model (torch.nn.Module): The PyTorch model to export.
im (torch.Tensor): Example input tensor for the model.
onnx_file (str): Path to save the exported ONNX file.
model (torch.nn.Module): The PyTorch model to export.
im (torch.Tensor | tuple[torch.Tensor, ...]): Example input tensor(s) for tracing.
output_file (Path | str): Path to save the exported ONNX file.
opset (int): ONNX opset version to use for export.
input_names (list[str]): List of input tensor names.
output_names (list[str]): List of output tensor names.
dynamic (bool | dict, optional): Whether to enable dynamic axes.
input_names (list[str] | None): List of input tensor names. Defaults to ``["images"]``.
output_names (list[str] | None): List of output tensor names. Defaults to ``["output0"]``.
dynamic (dict | None): Dictionary specifying dynamic axes for inputs and outputs.
Returns:
(str): Path to the exported ONNX file.
Notes:
Setting `do_constant_folding=True` may cause issues with DNN inference for torch>=1.12.
"""
if input_names is None:
input_names = ["images"]
if output_names is None:
output_names = ["output0"]
kwargs = {"dynamo": False} if TORCH_2_4 else {}
torch.onnx.export(
torch_model,
model,
im,
onnx_file,
output_file,
verbose=False,
opset_version=opset,
do_constant_folding=True, # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
input_names=input_names,
output_names=output_names,
dynamic_axes=dynamic or None,
dynamic_axes=dynamic,
**kwargs,
)
return str(output_file)
def onnx2engine(
onnx_file: str,
engine_file: str | None = None,
output_file: Path | str | None = None,
workspace: int | None = None,
half: bool = False,
int8: bool = False,
@ -94,12 +102,12 @@ def onnx2engine(
metadata: dict | None = None,
verbose: bool = False,
prefix: str = "",
) -> None:
) -> str:
"""Export a YOLO model to TensorRT engine format.
Args:
onnx_file (str): Path to the ONNX file to be converted.
engine_file (str | None): Path to save the generated TensorRT engine file.
output_file (Path | str | None): Path to save the generated TensorRT engine file.
workspace (int | None): Workspace size in GB for TensorRT.
half (bool, optional): Enable FP16 precision.
int8 (bool, optional): Enable INT8 precision.
@ -111,6 +119,9 @@ def onnx2engine(
verbose (bool, optional): Enable verbose logging.
prefix (str, optional): Prefix for log messages.
Returns:
(str): Path to the exported engine file.
Raises:
ValueError: If DLA is enabled on non-Jetson devices or required precision is not set.
RuntimeError: If the ONNX file cannot be parsed.
@ -122,7 +133,7 @@ def onnx2engine(
"""
import tensorrt as trt
engine_file = engine_file or Path(onnx_file).with_suffix(".engine")
output_file = output_file or Path(onnx_file).with_suffix(".engine")
logger = trt.Logger(trt.Logger.INFO)
if verbose:
@ -178,7 +189,7 @@ def onnx2engine(
if int8 and not is_trt10: # deprecated in TensorRT 10, causes internal errors
config.set_calibration_profile(profile)
LOGGER.info(f"{prefix} building {'INT8' if int8 else 'FP' + ('16' if half else '32')} engine as {engine_file}")
LOGGER.info(f"{prefix} building {'INT8' if int8 else 'FP' + ('16' if half else '32')} engine as {output_file}")
if int8:
config.set_flag(trt.BuilderFlag.INT8)
config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
@ -263,16 +274,17 @@ def onnx2engine(
engine = builder.build_serialized_network(network, config)
if engine is None:
raise RuntimeError("TensorRT engine build failed, check logs for errors")
with open(engine_file, "wb") as t:
with open(output_file, "wb") as t:
if metadata is not None:
meta = json.dumps(metadata)
t.write(len(meta).to_bytes(4, byteorder="little", signed=True))
t.write(meta.encode())
t.write(engine)
else:
with builder.build_engine(network, config) as engine, open(engine_file, "wb") as t:
with builder.build_engine(network, config) as engine, open(output_file, "wb") as t:
if metadata is not None:
meta = json.dumps(metadata)
t.write(len(meta).to_bytes(4, byteorder="little", signed=True))
t.write(meta.encode())
t.write(engine.serialize())
return str(output_file)

15
ultralytics/utils/export/executorch.py
View file

@ -39,8 +39,8 @@ def _executorch_kpts_decode(self, kpts: torch.Tensor, is_pose26: bool = False) -
def torch2executorch(
model: torch.nn.Module,
file: Path | str,
sample_input: torch.Tensor,
im: torch.Tensor,
output_dir: Path | str,
metadata: dict | None = None,
prefix: str = "",
) -> str:
@ -48,8 +48,8 @@ def torch2executorch(
Args:
model (torch.nn.Module): The PyTorch model to export.
file (Path | str): Source model file path used to derive output names.
sample_input (torch.Tensor): Example input tensor for tracing/export.
im (torch.Tensor): Example input tensor for tracing/export.
output_dir (Path | str): Directory to save the exported ExecuTorch model.
metadata (dict | None, optional): Optional metadata to save as YAML.
prefix (str, optional): Prefix for log messages.
@ -62,13 +62,12 @@ def torch2executorch(
LOGGER.info(f"\n{prefix} starting export with ExecuTorch {executorch_version.__version__}...")
file = Path(file)
output_dir = Path(str(file).replace(file.suffix, "_executorch_model"))
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
pte_file = output_dir / file.with_suffix(".pte").name
pte_file = output_dir / "model.pte"
et_program = to_edge_transform_and_lower(
torch.export.export(model, (sample_input,)),
torch.export.export(model, (im,)),
partitioner=[XnnpackPartitioner()],
).to_executorch()
pte_file.write_bytes(et_program.buffer)

26
ultralytics/utils/export/imx.py
View file

@ -203,7 +203,7 @@ class NMSWrapper(torch.nn.Module):
def torch2imx(
model: torch.nn.Module,
file: Path | str,
output_dir: Path | str,
conf: float,
iou: float,
max_det: int,
@ -211,7 +211,7 @@ def torch2imx(
gptq: bool = False,
dataset=None,
prefix: str = "",
):
) -> str:
"""Export YOLO model to IMX format for deployment on Sony IMX500 devices.
This function quantizes a YOLO model using Model Compression Toolkit (MCT) and exports it to IMX format compatible
@ -220,7 +220,7 @@ def torch2imx(
Args:
model (torch.nn.Module): The YOLO model to export. Must be YOLOv8n or YOLO11n.
file (Path | str): Output file path for the exported model.
output_dir (Path | str): Directory to save the exported IMX model.
conf (float): Confidence threshold for NMS post-processing.
iou (float): IoU threshold for NMS post-processing.
max_det (int): Maximum number of detections to return.
@ -231,7 +231,7 @@ def torch2imx(
prefix (str, optional): Logging prefix string. Defaults to "".
Returns:
(Path): Path to the exported IMX model directory.
(str): Path to the exported IMX model directory.
Raises:
ValueError: If the model is not a supported YOLOv8n or YOLO11n variant.
@ -239,7 +239,7 @@ def torch2imx(
Examples:
>>> from ultralytics import YOLO
>>> model = YOLO("yolo11n.pt")
>>> path = torch2imx(model, "model.imx", conf=0.25, iou=0.7, max_det=300)
>>> path = torch2imx(model, "output_dir/", conf=0.25, iou=0.7, max_det=300)
Notes:
- Requires model_compression_toolkit, onnx, edgemdt_tpc, and edge-mdt-cl packages
@ -309,9 +309,9 @@ def torch2imx(
task=model.task,
)
f = Path(str(file).replace(file.suffix, "_imx_model"))
f.mkdir(exist_ok=True)
onnx_model = f / Path(str(file.name).replace(file.suffix, "_imx.onnx")) # js dir
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
onnx_model = output_dir / "model_imx.onnx"
with onnx_export_patch():
mct.exporter.pytorch_export_model(
@ -319,7 +319,7 @@ def torch2imx(
)
model_onnx = onnx.load(onnx_model) # load onnx model
for k, v in metadata.items():
for k, v in (metadata or {}).items():
meta = model_onnx.metadata_props.add()
meta.key, meta.value = k, str(v)
@ -334,12 +334,12 @@ def torch2imx(
raise FileNotFoundError("imxconv-pt not found. Install with: pip install imx500-converter[pt]")
subprocess.run(
[str(imxconv), "-i", str(onnx_model), "-o", str(f), "--no-input-persistency", "--overwrite-output"],
[str(imxconv), "-i", str(onnx_model), "-o", str(output_dir), "--no-input-persistency", "--overwrite-output"],
check=True,
)
# Needed for imx models.
with open(f / "labels.txt", "w", encoding="utf-8") as file:
file.writelines([f"{name}\n" for _, name in model.names.items()])
with open(output_dir / "labels.txt", "w", encoding="utf-8") as labels_file:
labels_file.writelines([f"{name}\n" for _, name in model.names.items()])
return f
return str(output_dir)

28
ultralytics/utils/export/mnn.py
View file

@ -9,8 +9,8 @@ from ultralytics.utils import LOGGER
def onnx2mnn(
f_onnx: str,
file: Path | str,
onnx_file: str,
output_file: Path | str,
half: bool = False,
int8: bool = False,
metadata: dict | None = None,
@ -19,8 +19,8 @@ def onnx2mnn(
"""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.
onnx_file (str): Path to the source ONNX file.
output_file (Path | str): Path to save the exported MNN model.
half (bool): Whether to enable FP16 conversion.
int8 (bool): Whether to enable INT8 weight quantization.
metadata (dict | None): Optional metadata embedded via ``--bizCode``.
@ -33,23 +33,31 @@ def onnx2mnn(
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}"
assert Path(onnx_file).exists(), f"failed to export ONNX file: {onnx_file}"
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 {})]
mnn_args = [
"",
"-f",
"ONNX",
"--modelFile",
onnx_file,
"--MNNModel",
str(output_file),
"--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"
convert_scratch = Path(output_file).parent / ".__convert_external_data.bin"
if convert_scratch.exists():
convert_scratch.unlink()
return f
return str(output_file)

30
ultralytics/utils/export/ncnn.py
View file

@ -2,7 +2,6 @@
from __future__ import annotations
import os
from pathlib import Path
import torch
@ -13,7 +12,7 @@ from ultralytics.utils import LOGGER, YAML
def torch2ncnn(
model: torch.nn.Module,
im: torch.Tensor,
file: Path | str,
output_dir: Path | str,
half: bool = False,
metadata: dict | None = None,
device: torch.device | None = None,
@ -24,7 +23,7 @@ def torch2ncnn(
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.
output_dir (Path | str): Directory to save the exported NCNN model.
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.
@ -41,23 +40,22 @@ def torch2ncnn(
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}"))
output_dir = Path(output_dir)
ncnn_args = dict(
ncnnparam=(f / "model.ncnn.param").as_posix(),
ncnnbin=(f / "model.ncnn.bin").as_posix(),
ncnnpy=(f / "model_ncnn.py").as_posix(),
ncnnparam=(output_dir / "model.ncnn.param").as_posix(),
ncnnbin=(output_dir / "model.ncnn.bin").as_posix(),
ncnnpy=(output_dir / "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(),
ptpath=(output_dir / "model.pt").as_posix(),
pnnxparam=(output_dir / "model.pnnx.param").as_posix(),
pnnxbin=(output_dir / "model.pnnx.bin").as_posix(),
pnnxpy=(output_dir / "model_pnnx.py").as_posix(),
pnnxonnx=(output_dir / "model.pnnx.onnx").as_posix(),
)
f.mkdir(exist_ok=True) # make ncnn_model directory
output_dir.mkdir(parents=True, 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)
@ -65,5 +63,5 @@ def torch2ncnn(
Path(f_debug).unlink(missing_ok=True)
if metadata:
YAML.save(f / "metadata.yaml", metadata) # add metadata.yaml
return str(f)
YAML.save(output_dir / "metadata.yaml", metadata) # add metadata.yaml
return str(output_dir)

27
ultralytics/utils/export/openvino.py
View file

@ -2,7 +2,6 @@
from __future__ import annotations
import os
from pathlib import Path
from typing import Any
@ -13,25 +12,25 @@ from ultralytics.utils import LOGGER
def torch2openvino(
model: torch.nn.Module,
im: torch.Tensor,
file: Path | str | None = None,
im: torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor, ...],
output_dir: 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:
) -> Any:
"""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.
im (torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor, ...]): Example input tensor(s) for tracing.
output_dir (Path | str | None): Directory to save the exported OpenVINO model.
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``).
calibration_dataset (nncf.Dataset | None): Dataset for INT8 calibration (required when ``int8=True``).
ignored_scope (dict | None): Kwargs passed to ``nncf.IgnoredScope`` for head patterns.
prefix (str): Prefix for log messages.
@ -42,7 +41,8 @@ def torch2openvino(
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)
input_shape = [i.shape for i in im] if isinstance(im, (list, tuple)) else im.shape
ov_model = ov.convert_model(model, input=None if dynamic else input_shape, example_input=im)
if int8:
import nncf
@ -53,10 +53,9 @@ def torch2openvino(
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)
if output_dir is not None:
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
output_file = output_dir / "model.xml"
ov.save_model(ov_model, output_file, compress_to_fp16=half)
return ov_model

13
ultralytics/utils/export/paddle.py
View file

@ -2,7 +2,6 @@
from __future__ import annotations
import os
from pathlib import Path
import torch
@ -13,7 +12,7 @@ from ultralytics.utils import ARM64, IS_JETSON, LOGGER, YAML
def torch2paddle(
model: torch.nn.Module,
im: torch.Tensor,
file: Path | str,
output_dir: Path | str,
metadata: dict | None = None,
prefix: str = "",
) -> str:
@ -22,7 +21,7 @@ def torch2paddle(
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.
output_dir (Path | str): Directory to save the exported PaddlePaddle model.
metadata (dict | None): Optional metadata saved as ``metadata.yaml``.
prefix (str): Prefix for log messages.
@ -47,10 +46,8 @@ def torch2paddle(
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
pytorch2paddle(module=model, save_dir=output_dir, jit_type="trace", input_examples=[im]) # export
if metadata:
YAML.save(Path(f) / "metadata.yaml", metadata) # add metadata.yaml
return f
YAML.save(Path(output_dir) / "metadata.yaml", metadata) # add metadata.yaml
return str(output_dir)

22
ultralytics/utils/export/rknn.py
View file

@ -8,21 +8,23 @@ from ultralytics.utils import IS_COLAB, LOGGER, YAML
def onnx2rknn(
f_onnx: str,
onnx_file: str,
output_dir: Path | str,
name: str = "rk3588",
metadata: dict | None = None,
prefix: str = "",
) -> Path:
) -> str:
"""Export an ONNX model to RKNN format for Rockchip NPUs.
Args:
f_onnx (str): Path to the source ONNX file (already exported, opset <=19).
onnx_file (str): Path to the source ONNX file (already exported, opset <=19).
output_dir (Path | str): Directory to save the exported RKNN model.
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.
(str): Path to the exported ``_rknn_model`` directory.
"""
from ultralytics.utils.checks import check_requirements
@ -38,14 +40,14 @@ def onnx2rknn(
from rknn.api import RKNN
export_path = Path(f"{Path(f_onnx).stem}_rknn_model")
export_path.mkdir(exist_ok=True)
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, 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.load_onnx(model=onnx_file)
rknn.build(do_quantization=False) # TODO: Add quantization support
rknn.export_rknn(str(export_path / f"{Path(f_onnx).stem}-{name}.rknn"))
rknn.export_rknn(str(output_dir / f"{Path(onnx_file).stem}-{name}.rknn"))
if metadata:
YAML.save(export_path / "metadata.yaml", metadata)
return export_path
YAML.save(output_dir / "metadata.yaml", metadata)
return str(output_dir)

42
ultralytics/utils/export/tensorflow.py
View file

@ -59,19 +59,19 @@ def _tf_kpts_decode(self, kpts: torch.Tensor, is_pose26: bool = False) -> torch.
def onnx2saved_model(
onnx_file: str,
output_dir: Path,
output_dir: Path | str,
int8: bool = False,
images: np.ndarray = None,
images: np.ndarray | None = None,
disable_group_convolution: bool = False,
prefix="",
prefix: str = "",
):
"""Convert an ONNX model to TensorFlow SavedModel format using onnx2tf.
Args:
onnx_file (str): ONNX file path.
output_dir (Path): Output directory path for the SavedModel.
output_dir (Path | str): Output directory path for the SavedModel.
int8 (bool, optional): Enable INT8 quantization. Defaults to False.
images (np.ndarray, optional): Calibration images for INT8 quantization in BHWC format.
images (np.ndarray | None, optional): Calibration images for INT8 quantization in BHWC format.
disable_group_convolution (bool, optional): Disable group convolution optimization. Defaults to False.
prefix (str, optional): Logging prefix. Defaults to "".
@ -82,6 +82,7 @@ def onnx2saved_model(
- Requires onnx2tf package. Downloads calibration data if INT8 quantization is enabled.
- Removes temporary files and renames quantized models after conversion.
"""
output_dir = Path(output_dir)
# Pre-download calibration file to fix https://github.com/PINTO0309/onnx2tf/issues/545
onnx2tf_file = Path("calibration_image_sample_data_20x128x128x3_float32.npy")
if not onnx2tf_file.exists():
@ -118,7 +119,7 @@ def onnx2saved_model(
verbosity="error", # note INT8-FP16 activation bug https://github.com/ultralytics/ultralytics/issues/15873
output_integer_quantized_tflite=int8,
custom_input_op_name_np_data_path=np_data,
enable_batchmatmul_unfold=True and not int8, # fix lower no. of detected objects on GPU delegate
enable_batchmatmul_unfold=not int8, # fix lower no. of detected objects on GPU delegate
output_signaturedefs=True, # fix error with Attention block group convolution
disable_group_convolution=disable_group_convolution, # fix error with group convolution
)
@ -133,14 +134,17 @@ def onnx2saved_model(
return keras_model
def keras2pb(keras_model, file: Path, prefix=""):
def keras2pb(keras_model, output_file: Path | str, prefix: str = "") -> str:
"""Convert a Keras model to TensorFlow GraphDef (.pb) format.
Args:
keras_model (keras.Model): Keras model to convert to frozen graph format.
file (Path): Output file path (suffix will be changed to .pb).
output_file (Path | str): Output file path (suffix will be changed to .pb).
prefix (str, optional): Logging prefix. Defaults to "".
Returns:
(str): Path to the exported ``.pb`` file.
Notes:
Creates a frozen graph by converting variables to constants for inference optimization.
"""
@ -152,10 +156,14 @@ def keras2pb(keras_model, file: Path, prefix=""):
m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
frozen_func = convert_variables_to_constants_v2(m)
frozen_func.graph.as_graph_def()
tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(file.parent), name=file.name, as_text=False)
output_file = Path(output_file)
tf.io.write_graph(
graph_or_graph_def=frozen_func.graph, logdir=str(output_file.parent), name=output_file.name, as_text=False
)
return str(output_file)
def tflite2edgetpu(tflite_file: str | Path, output_dir: str | Path, prefix: str = ""):
def tflite2edgetpu(tflite_file: str | Path, output_dir: str | Path, prefix: str = "") -> str:
"""Convert a TensorFlow Lite model to Edge TPU format using the Edge TPU compiler.
Args:
@ -163,6 +171,9 @@ def tflite2edgetpu(tflite_file: str | Path, output_dir: str | Path, prefix: str
output_dir (str | Path): Output directory path for the compiled Edge TPU model.
prefix (str, optional): Logging prefix. Defaults to "".
Returns:
(str): Path to the exported Edge TPU model file.
Notes:
Requires the Edge TPU compiler to be installed. The function compiles the TFLite model
for optimal performance on Google's Edge TPU hardware accelerator.
@ -180,9 +191,10 @@ def tflite2edgetpu(tflite_file: str | Path, output_dir: str | Path, prefix: str
)
LOGGER.info(f"{prefix} running '{cmd}'")
subprocess.run(cmd, shell=True)
return str(Path(output_dir) / f"{Path(tflite_file).stem}_edgetpu.tflite")
def pb2tfjs(pb_file: str, output_dir: str, half: bool = False, int8: bool = False, prefix: str = ""):
def pb2tfjs(pb_file: str, output_dir: str, half: bool = False, int8: bool = False, prefix: str = "") -> str:
"""Convert a TensorFlow GraphDef (.pb) model to TensorFlow.js format.
Args:
@ -192,6 +204,9 @@ def pb2tfjs(pb_file: str, output_dir: str, half: bool = False, int8: bool = Fals
int8 (bool, optional): Enable INT8 quantization. Defaults to False.
prefix (str, optional): Logging prefix. Defaults to "".
Returns:
(str): Path to the exported TensorFlow.js model directory.
Notes:
Requires tensorflowjs package. Uses tensorflowjs_converter command-line tool for conversion.
Handles spaces in file paths and warns if output directory contains spaces.
@ -204,8 +219,8 @@ def pb2tfjs(pb_file: str, output_dir: str, half: bool = False, int8: bool = Fals
LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...")
gd = tf.Graph().as_graph_def() # TF GraphDef
with open(pb_file, "rb") as file:
gd.ParseFromString(file.read())
with open(pb_file, "rb") as f:
gd.ParseFromString(f.read())
outputs = ",".join(gd_outputs(gd))
LOGGER.info(f"\n{prefix} output node names: {outputs}")
@ -220,6 +235,7 @@ def pb2tfjs(pb_file: str, output_dir: str, half: bool = False, int8: bool = Fals
if " " in output_dir:
LOGGER.warning(f"{prefix} your model may not work correctly with spaces in path '{output_dir}'.")
return str(output_dir)
def gd_outputs(gd):

17
ultralytics/utils/export/torchscript.py
View file

@ -13,35 +13,34 @@ from ultralytics.utils import LOGGER, TORCH_VERSION
def torch2torchscript(
model: torch.nn.Module,
im: torch.Tensor,
file: Path | str,
output_file: Path | str,
optimize: bool = False,
metadata: dict | None = None,
prefix: str = "",
) -> Path:
) -> str:
"""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.
output_file (Path | str): Path to save the exported TorchScript model.
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.
(str): 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")
output_file = str(output_file)
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)
optimize_for_mobile(ts)._save_for_lite_interpreter(output_file, _extra_files=extra_files)
else:
ts.save(str(f), _extra_files=extra_files)
return f
ts.save(output_file, _extra_files=extra_files)
return output_file