diff --git a/docs/en/integrations/sony-imx500.md b/docs/en/integrations/sony-imx500.md index 5d2eed7f96..40a65ab822 100644 --- a/docs/en/integrations/sony-imx500.md +++ b/docs/en/integrations/sony-imx500.md @@ -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 diff --git a/tests/test_exports.py b/tests/test_exports.py index 38fcd668ba..51b73040d5 100644 --- a/tests/test_exports.py +++ b/tests/test_exports.py @@ -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" diff --git a/ultralytics/__init__.py b/ultralytics/__init__.py index a939faa1cb..1184e63c07 100644 --- a/ultralytics/__init__.py +++ b/ultralytics/__init__.py @@ -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 diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py index 4565f86f46..32ee9d544f 100644 --- a/ultralytics/engine/exporter.py +++ b/ultralytics/engine/exporter.py @@ -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, diff --git a/ultralytics/nn/backends/coreml.py b/ultralytics/nn/backends/coreml.py index 9f96c66ec7..d36dfc2e56 100644 --- a/ultralytics/nn/backends/coreml.py +++ b/ultralytics/nn/backends/coreml.py @@ -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 diff --git a/ultralytics/utils/export/axelera.py b/ultralytics/utils/export/axelera.py index cdd3b0f390..35cd24f516 100644 --- a/ultralytics/utils/export/axelera.py +++ b/ultralytics/utils/export/axelera.py @@ -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) diff --git a/ultralytics/utils/export/coreml.py b/ultralytics/utils/export/coreml.py index 12c19dd129..6c6c72f15f 100644 --- a/ultralytics/utils/export/coreml.py +++ b/ultralytics/utils/export/coreml.py @@ -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 diff --git a/ultralytics/utils/export/engine.py b/ultralytics/utils/export/engine.py index a2ce93e458..ab619761c4 100644 --- a/ultralytics/utils/export/engine.py +++ b/ultralytics/utils/export/engine.py @@ -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) diff --git a/ultralytics/utils/export/executorch.py b/ultralytics/utils/export/executorch.py index 5606be8acc..15e00805b9 100644 --- a/ultralytics/utils/export/executorch.py +++ b/ultralytics/utils/export/executorch.py @@ -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) diff --git a/ultralytics/utils/export/imx.py b/ultralytics/utils/export/imx.py index ba8e462d71..21689891fa 100644 --- a/ultralytics/utils/export/imx.py +++ b/ultralytics/utils/export/imx.py @@ -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) diff --git a/ultralytics/utils/export/mnn.py b/ultralytics/utils/export/mnn.py index 3cebf890ec..0142f83385 100644 --- a/ultralytics/utils/export/mnn.py +++ b/ultralytics/utils/export/mnn.py @@ -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) diff --git a/ultralytics/utils/export/ncnn.py b/ultralytics/utils/export/ncnn.py index 28d33517d1..177e267ee7 100644 --- a/ultralytics/utils/export/ncnn.py +++ b/ultralytics/utils/export/ncnn.py @@ -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) diff --git a/ultralytics/utils/export/openvino.py b/ultralytics/utils/export/openvino.py index c804a982d5..d15f07dbda 100644 --- a/ultralytics/utils/export/openvino.py +++ b/ultralytics/utils/export/openvino.py @@ -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 diff --git a/ultralytics/utils/export/paddle.py b/ultralytics/utils/export/paddle.py index 841a01a919..7d3f467019 100644 --- a/ultralytics/utils/export/paddle.py +++ b/ultralytics/utils/export/paddle.py @@ -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) diff --git a/ultralytics/utils/export/rknn.py b/ultralytics/utils/export/rknn.py index d1aacec58b..1c8da6ba9d 100644 --- a/ultralytics/utils/export/rknn.py +++ b/ultralytics/utils/export/rknn.py @@ -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) diff --git a/ultralytics/utils/export/tensorflow.py b/ultralytics/utils/export/tensorflow.py index c70a709875..322edb76c9 100644 --- a/ultralytics/utils/export/tensorflow.py +++ b/ultralytics/utils/export/tensorflow.py @@ -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): diff --git a/ultralytics/utils/export/torchscript.py b/ultralytics/utils/export/torchscript.py index 6805eac217..077b69f075 100644 --- a/ultralytics/utils/export/torchscript.py +++ b/ultralytics/utils/export/torchscript.py @@ -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