ultralytics 8.3.143__py3-none-any.whl → 8.3.145__py3-none-any.whl

ultralytics/engine/exporter.py

@@ -222,11 +222,53 @@ def arange_patch(args):
 
 class Exporter:
     """
-    A class for exporting a model.
+    A class for exporting YOLO models to various formats.
+
+    This class provides functionality to export YOLO models to different formats including ONNX, TensorRT, CoreML,
+    TensorFlow, and others. It handles format validation, device selection, model preparation, and the actual export
+    process for each supported format.
 
     Attributes:
-        args (SimpleNamespace): Configuration for the exporter.
-        callbacks (list, optional): List of callback functions.
+        args (SimpleNamespace): Configuration arguments for the exporter.
+        callbacks (dict): Dictionary of callback functions for different export events.
+        im (torch.Tensor): Input tensor for model inference during export.
+        model (torch.nn.Module): The YOLO model to be exported.
+        file (Path): Path to the model file being exported.
+        output_shape (tuple): Shape of the model output tensor(s).
+        pretty_name (str): Formatted model name for display purposes.
+        metadata (dict): Model metadata including description, author, version, etc.
+        device (torch.device): Device on which the model is loaded.
+        imgsz (tuple): Input image size for the model.
+
+    Methods:
+        __call__: Main export method that handles the export process.
+        get_int8_calibration_dataloader: Build dataloader for INT8 calibration.
+        export_torchscript: Export model to TorchScript format.
+        export_onnx: Export model to ONNX format.
+        export_openvino: Export model to OpenVINO format.
+        export_paddle: Export model to PaddlePaddle format.
+        export_mnn: Export model to MNN format.
+        export_ncnn: Export model to NCNN format.
+        export_coreml: Export model to CoreML format.
+        export_engine: Export model to TensorRT format.
+        export_saved_model: Export model to TensorFlow SavedModel format.
+        export_pb: Export model to TensorFlow GraphDef format.
+        export_tflite: Export model to TensorFlow Lite format.
+        export_edgetpu: Export model to Edge TPU format.
+        export_tfjs: Export model to TensorFlow.js format.
+        export_rknn: Export model to RKNN format.
+        export_imx: Export model to IMX format.
+
+    Examples:
+        Export a YOLOv8 model to ONNX format
+        >>> from ultralytics.engine.exporter import Exporter
+        >>> exporter = Exporter()
+        >>> exporter(model="yolov8n.pt")  # exports to yolov8n.onnx
+
+        Export with specific arguments
+        >>> args = {"format": "onnx", "dynamic": True, "half": True}
+        >>> exporter = Exporter(overrides=args)
+        >>> exporter(model="yolov8n.pt")
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
@@ -536,7 +578,7 @@ class Exporter:
 
     @try_export
     def export_torchscript(self, prefix=colorstr("TorchScript:")):
-        """YOLO TorchScript model export."""
+        """Export YOLO model to TorchScript format."""
         LOGGER.info(f"\n{prefix} starting export with torch {torch.__version__}...")
         f = self.file.with_suffix(".torchscript")
 
@@ -553,7 +595,7 @@ class Exporter:
 
     @try_export
     def export_onnx(self, prefix=colorstr("ONNX:")):
-        """YOLO ONNX export."""
+        """Export YOLO model to ONNX format."""
         requirements = ["onnx>=1.12.0,<1.18.0"]
         if self.args.simplify:
             requirements += ["onnxslim>=0.1.53", "onnxruntime" + ("-gpu" if torch.cuda.is_available() else "")]
@@ -612,7 +654,7 @@ class Exporter:
 
     @try_export
     def export_openvino(self, prefix=colorstr("OpenVINO:")):
-        """YOLO OpenVINO export."""
+        """Export YOLO model to OpenVINO format."""
         if MACOS:
             msg = "OpenVINO error in macOS>=15.4 https://github.com/openvinotoolkit/openvino/issues/30023"
             check_version(MACOS_VERSION, "<15.4", name="macOS ", hard=True, msg=msg)
@@ -689,7 +731,7 @@ class Exporter:
 
     @try_export
     def export_paddle(self, prefix=colorstr("PaddlePaddle:")):
-        """YOLO Paddle export."""
+        """Export YOLO model to PaddlePaddle format."""
         assert not IS_JETSON, "Jetson Paddle exports not supported yet"
         check_requirements(("paddlepaddle-gpu" if torch.cuda.is_available() else "paddlepaddle>=3.0.0", "x2paddle"))
         import x2paddle  # noqa
@@ -704,7 +746,7 @@ class Exporter:
 
     @try_export
     def export_mnn(self, prefix=colorstr("MNN:")):
-        """YOLO MNN export using MNN https://github.com/alibaba/MNN."""
+        """Export YOLO model to MNN format using MNN https://github.com/alibaba/MNN."""
         f_onnx, _ = self.export_onnx()  # get onnx model first
 
         check_requirements("MNN>=2.9.6")
@@ -729,7 +771,7 @@ class Exporter:
 
     @try_export
     def export_ncnn(self, prefix=colorstr("NCNN:")):
-        """YOLO NCNN export using PNNX https://github.com/pnnx/pnnx."""
+        """Export YOLO model to NCNN format using PNNX https://github.com/pnnx/pnnx."""
         check_requirements("ncnn")
         import ncnn  # noqa
 
@@ -797,7 +839,7 @@ class Exporter:
 
     @try_export
     def export_coreml(self, prefix=colorstr("CoreML:")):
-        """YOLO CoreML export."""
+        """Export YOLO model to CoreML format."""
         mlmodel = self.args.format.lower() == "mlmodel"  # legacy *.mlmodel export format requested
         check_requirements("coremltools>=8.0")
         import coremltools as ct  # noqa
@@ -876,7 +918,7 @@ class Exporter:
 
     @try_export
     def export_engine(self, dla=None, prefix=colorstr("TensorRT:")):
-        """YOLO TensorRT export https://developer.nvidia.com/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
 
@@ -912,7 +954,7 @@ class Exporter:
 
     @try_export
     def export_saved_model(self, prefix=colorstr("TensorFlow SavedModel:")):
-        """YOLO TensorFlow SavedModel export."""
+        """Export YOLO model to TensorFlow SavedModel format."""
         cuda = torch.cuda.is_available()
         try:
             import tensorflow as tf  # noqa
@@ -1002,7 +1044,7 @@ class Exporter:
 
     @try_export
     def export_pb(self, keras_model, prefix=colorstr("TensorFlow GraphDef:")):
-        """YOLO TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen-Graph-TensorFlow."""
+        """Export YOLO model to TensorFlow GraphDef *.pb format https://github.com/leimao/Frozen-Graph-TensorFlow."""
         import tensorflow as tf  # noqa
         from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2  # noqa
 
@@ -1018,7 +1060,7 @@ class Exporter:
 
     @try_export
     def export_tflite(self, prefix=colorstr("TensorFlow Lite:")):
-        """YOLO TensorFlow Lite export."""
+        """Export YOLO model to TensorFlow Lite format."""
         # BUG https://github.com/ultralytics/ultralytics/issues/13436
         import tensorflow as tf  # noqa
 
@@ -1034,7 +1076,7 @@ class Exporter:
 
     @try_export
     def export_edgetpu(self, tflite_model="", prefix=colorstr("Edge TPU:")):
-        """YOLO Edge TPU export https://coral.ai/docs/edgetpu/models-intro/."""
+        """Export YOLO model to Edge TPU format https://coral.ai/docs/edgetpu/models-intro/."""
         cmd = "edgetpu_compiler --version"
         help_url = "https://coral.ai/docs/edgetpu/compiler/"
         assert LINUX, f"export only supported on Linux. See {help_url}"
@@ -1069,7 +1111,7 @@ class Exporter:
 
     @try_export
     def export_tfjs(self, prefix=colorstr("TensorFlow.js:")):
-        """YOLO TensorFlow.js export."""
+        """Export YOLO model to TensorFlow.js format."""
         check_requirements("tensorflowjs")
         import tensorflow as tf
         import tensorflowjs as tfjs  # noqa
@@ -1102,7 +1144,7 @@ class Exporter:
 
     @try_export
     def export_rknn(self, prefix=colorstr("RKNN:")):
-        """YOLO RKNN model export."""
+        """Export YOLO model to RKNN format."""
         LOGGER.info(f"\n{prefix} starting export with rknn-toolkit2...")
 
         check_requirements("rknn-toolkit2")
@@ -1129,7 +1171,7 @@ class Exporter:
 
     @try_export
     def export_imx(self, prefix=colorstr("IMX:")):
-        """YOLO IMX export."""
+        """Export YOLO model to IMX format."""
         gptq = False
         assert LINUX, (
             "export only supported on Linux. "
@@ -1212,6 +1254,8 @@ class Exporter:
         )
 
         class NMSWrapper(torch.nn.Module):
+            """Wrap PyTorch Module with multiclass_nms layer from sony_custom_layers."""
+
             def __init__(
                 self,
                 model: torch.nn.Module,
@@ -1220,13 +1264,13 @@ class Exporter:
                 max_detections: int = 300,
             ):
                 """
-                Wrapping PyTorch Module with multiclass_nms layer from sony_custom_layers.
+                Initialize NMSWrapper with PyTorch Module and NMS parameters.
 
                 Args:
-                    model (nn.Module): Model instance.
+                    model (torch.nn.Module): Model instance.
                     score_threshold (float): Score threshold for non-maximum suppression.
                     iou_threshold (float): Intersection over union threshold for non-maximum suppression.
-                    max_detections (float): The number of detections to return.
+                    max_detections (int): The number of detections to return.
                 """
                 super().__init__()
                 self.model = model
@@ -1235,6 +1279,7 @@ class Exporter:
                 self.max_detections = max_detections
 
             def forward(self, images):
+                """Forward pass with model inference and NMS post-processing."""
                 # model inference
                 outputs = self.model(images)
 
@@ -1289,7 +1334,7 @@ class Exporter:
             zf.writestr("metadata.json", json.dumps(self.metadata, indent=2))
 
     def _pipeline_coreml(self, model, weights_dir=None, prefix=colorstr("CoreML Pipeline:")):
-        """YOLO CoreML pipeline."""
+        """Create CoreML pipeline with NMS for YOLO detection models."""
         import coremltools as ct  # noqa
 
         LOGGER.info(f"{prefix} starting pipeline with coremltools {ct.__version__}...")
@@ -1395,7 +1440,7 @@ class Exporter:
         return model
 
     def add_callback(self, event: str, callback):
-        """Appends the given callback."""
+        """Append the given callback to the specified event."""
         self.callbacks[event].append(callback)
 
     def run_callbacks(self, event: str):
@@ -1408,7 +1453,13 @@ class IOSDetectModel(torch.nn.Module):
     """Wrap an Ultralytics YOLO model for Apple iOS CoreML export."""
 
     def __init__(self, model, im):
-        """Initialize the IOSDetectModel class with a YOLO model and example image."""
+        """
+        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).
+        """
         super().__init__()
         _, _, h, w = im.shape  # batch, channel, height, width
         self.model = model
@@ -1432,7 +1483,7 @@ class NMSModel(torch.nn.Module):
         Initialize the NMSModel.
 
         Args:
-            model (torch.nn.module): The model to wrap with NMS postprocessing.
+            model (torch.nn.Module): The model to wrap with NMS postprocessing.
             args (Namespace): The export arguments.
         """
         super().__init__()
@@ -1443,13 +1494,14 @@ class NMSModel(torch.nn.Module):
 
     def forward(self, x):
         """
-        Performs inference with NMS post-processing. Supports Detect, Segment, OBB and Pose.
+        Perform inference with NMS post-processing. Supports Detect, Segment, OBB and Pose.
 
         Args:
             x (torch.Tensor): The preprocessed tensor with shape (N, 3, H, W).
 
         Returns:
-            (torch.Tensor): List of detections, each an (N, max_det, 4 + 2 + extra_shape) Tensor where N is the number of detections after NMS.
+            (torch.Tensor): List of detections, each an (N, max_det, 4 + 2 + extra_shape) Tensor where N is the
+                number of detections after NMS.
         """
         from functools import partial
 

ultralytics/engine/model.py

@@ -48,25 +48,25 @@ class Model(torch.nn.Module):
 
     Methods:
         __call__: Alias for the predict method, enabling the model instance to be callable.
-        _new: Initializes a new model based on a configuration file.
-        _load: Loads a model from a checkpoint file.
-        _check_is_pytorch_model: Ensures that the model is a PyTorch model.
-        reset_weights: Resets the model's weights to their initial state.
-        load: Loads model weights from a specified file.
-        save: Saves the current state of the model to a file.
-        info: Logs or returns information about the model.
-        fuse: Fuses Conv2d and BatchNorm2d layers for optimized inference.
-        predict: Performs object detection predictions.
-        track: Performs object tracking.
-        val: Validates the model on a dataset.
-        benchmark: Benchmarks the model on various export formats.
-        export: Exports the model to different formats.
-        train: Trains the model on a dataset.
-        tune: Performs hyperparameter tuning.
-        _apply: Applies a function to the model's tensors.
-        add_callback: Adds a callback function for an event.
-        clear_callback: Clears all callbacks for an event.
-        reset_callbacks: Resets all callbacks to their default functions.
+        _new: Initialize a new model based on a configuration file.
+        _load: Load a model from a checkpoint file.
+        _check_is_pytorch_model: Ensure that the model is a PyTorch model.
+        reset_weights: Reset the model's weights to their initial state.
+        load: Load model weights from a specified file.
+        save: Save the current state of the model to a file.
+        info: Log or return information about the model.
+        fuse: Fuse Conv2d and BatchNorm2d layers for optimized inference.
+        predict: Perform object detection predictions.
+        track: Perform object tracking.
+        val: Validate the model on a dataset.
+        benchmark: Benchmark the model on various export formats.
+        export: Export the model to different formats.
+        train: Train the model on a dataset.
+        tune: Perform hyperparameter tuning.
+        _apply: Apply a function to the model's tensors.
+        add_callback: Add a callback function for an event.
+        clear_callback: Clear all callbacks for an event.
+        reset_callbacks: Reset all callbacks to their default functions.
 
     Examples:
         >>> from ultralytics import YOLO
@@ -94,7 +94,7 @@ class Model(torch.nn.Module):
         Args:
             model (str | Path | Model): Path or name of the model to load or create. Can be a local file path, a
                 model name from Ultralytics HUB, a Triton Server model, or an already initialized Model instance.
-            task (str | None): The task type associated with the YOLO model, specifying its application domain.
+            task (str, optional): The specific task for the model. If None, it will be inferred from the config.
             verbose (bool): If True, enables verbose output during the model's initialization and subsequent
                 operations.
 
@@ -242,9 +242,9 @@ class Model(torch.nn.Module):
 
         Args:
             cfg (str): Path to the model configuration file in YAML format.
-            task (str | None): The specific task for the model. If None, it will be inferred from the config.
-            model (torch.nn.Module | None): A custom model instance. If provided, it will be used instead of creating
-                a new one.
+            task (str, optional): The specific task for the model. If None, it will be inferred from the config.
+            model (torch.nn.Module, optional): A custom model instance. If provided, it will be used instead of
+                creating a new one.
             verbose (bool): If True, displays model information during loading.
 
         Raises:
@@ -276,7 +276,7 @@ class Model(torch.nn.Module):
 
         Args:
             weights (str): Path to the model weights file to be loaded.
-            task (str | None): The task associated with the model. If None, it will be inferred from the model.
+            task (str, optional): The task associated with the model. If None, it will be inferred from the model.
 
         Raises:
             FileNotFoundError: If the specified weights file does not exist or is inaccessible.
@@ -367,7 +367,7 @@ class Model(torch.nn.Module):
         name and shape and transfers them to the model.
 
         Args:
-            weights (Union[str, Path]): Path to the weights file or a weights object.
+            weights (str | Path): Path to the weights file or a weights object.
 
         Returns:
             (Model): The instance of the class with loaded weights.
@@ -501,7 +501,7 @@ class Model(torch.nn.Module):
         **kwargs: Any,
     ) -> List[Results]:
         """
-        Performs predictions on the given image source using the YOLO model.
+        Perform predictions on the given image source using the YOLO model.
 
         This method facilitates the prediction process, allowing various configurations through keyword arguments.
         It supports predictions with custom predictors or the default predictor method. The method handles different
@@ -512,7 +512,7 @@ class Model(torch.nn.Module):
                 of the image(s) to make predictions on. Accepts various types including file paths, URLs, PIL
                 images, numpy arrays, and torch tensors.
             stream (bool): If True, treats the input source as a continuous stream for predictions.
-            predictor (BasePredictor | None): An instance of a custom predictor class for making predictions.
+            predictor (BasePredictor, optional): An instance of a custom predictor class for making predictions.
                 If None, the method uses a default predictor.
             **kwargs (Any): Additional keyword arguments for configuring the prediction process.
 
@@ -562,14 +562,14 @@ class Model(torch.nn.Module):
         **kwargs: Any,
     ) -> List[Results]:
         """
-        Conducts object tracking on the specified input source using the registered trackers.
+        Conduct object tracking on the specified input source using the registered trackers.
 
         This method performs object tracking using the model's predictors and optionally registered trackers. It handles
         various input sources such as file paths or video streams, and supports customization through keyword arguments.
         The method registers trackers if not already present and can persist them between calls.
 
         Args:
-            source (Union[str, Path, int, List, Tuple, np.ndarray, torch.Tensor], optional): Input source for object
+            source (str | Path | int | List | Tuple | np.ndarray | torch.Tensor, optional): Input source for object
                 tracking. Can be a file path, URL, or video stream.
             stream (bool): If True, treats the input source as a continuous video stream.
             persist (bool): If True, persists trackers between different calls to this method.
@@ -611,8 +611,8 @@ class Model(torch.nn.Module):
         configurations, method-specific defaults, and user-provided arguments to configure the validation process.
 
         Args:
-            validator (ultralytics.engine.validator.BaseValidator | None): An instance of a custom validator class for
-                validating the model.
+            validator (ultralytics.engine.validator.BaseValidator, optional): An instance of a custom validator class
+                for validating the model.
             **kwargs (Any): Arbitrary keyword arguments for customizing the validation process.
 
         Returns:
@@ -634,10 +634,7 @@ class Model(torch.nn.Module):
         self.metrics = validator.metrics
         return validator.metrics
 
-    def benchmark(
-        self,
-        **kwargs: Any,
-    ):
+    def benchmark(self, data=None, format="", verbose=False, **kwargs: Any):
         """
         Benchmark the model across various export formats to evaluate performance.
 
@@ -647,14 +644,14 @@ class Model(torch.nn.Module):
         defaults, and any additional user-provided keyword arguments.
 
         Args:
+            data (str): Path to the dataset for benchmarking.
+            verbose (bool): Whether to print detailed benchmark information.
+            format (str): Export format name for specific benchmarking.
             **kwargs (Any): Arbitrary keyword arguments to customize the benchmarking process. Common options include:
-                - data (str): Path to the dataset for benchmarking.
                 - imgsz (int | List[int]): Image size for benchmarking.
                 - half (bool): Whether to use half-precision (FP16) mode.
                 - int8 (bool): Whether to use int8 precision mode.
                 - device (str): Device to run the benchmark on (e.g., 'cpu', 'cuda').
-                - verbose (bool): Whether to print detailed benchmark information.
-                - format (str): Export format name for specific benchmarking.
 
         Returns:
             (dict): A dictionary containing the results of the benchmarking process, including metrics for
@@ -671,17 +668,21 @@ class Model(torch.nn.Module):
         self._check_is_pytorch_model()
         from ultralytics.utils.benchmarks import benchmark
 
+        from .exporter import export_formats
+
         custom = {"verbose": False}  # method defaults
         args = {**DEFAULT_CFG_DICT, **self.model.args, **custom, **kwargs, "mode": "benchmark"}
+        fmts = export_formats()
+        export_args = set(dict(zip(fmts["Argument"], fmts["Arguments"])).get(format, []))
+        export_kwargs = {k: v for k, v in args.items() if k in export_args - set(["batch"])}
         return benchmark(
             model=self,
-            data=kwargs.get("data"),  # if no 'data' argument passed set data=None for default datasets
+            data=data,  # if no 'data' argument passed set data=None for default datasets
             imgsz=args["imgsz"],
-            half=args["half"],
-            int8=args["int8"],
             device=args["device"],
-            verbose=kwargs.get("verbose", False),
-            format=kwargs.get("format", ""),
+            verbose=verbose,
+            format=format,
+            **export_kwargs,
         )
 
     def export(
@@ -738,7 +739,7 @@ class Model(torch.nn.Module):
         **kwargs: Any,
     ):
         """
-        Trains the model using the specified dataset and training configuration.
+        Train the model using the specified dataset and training configuration.
 
         This method facilitates model training with a range of customizable settings. It supports training with a
         custom trainer or the default training approach. The method handles scenarios such as resuming training
@@ -749,7 +750,7 @@ class Model(torch.nn.Module):
         configurations, method-specific defaults, and user-provided arguments to configure the training process.
 
         Args:
-            trainer (BaseTrainer | None): Custom trainer instance for model training. If None, uses default.
+            trainer (BaseTrainer, optional): Custom trainer instance for model training. If None, uses default.
             **kwargs (Any): Arbitrary keyword arguments for training configuration. Common options include:
                 data (str): Path to dataset configuration file.
                 epochs (int): Number of training epochs.
@@ -810,7 +811,7 @@ class Model(torch.nn.Module):
         **kwargs: Any,
     ):
         """
-        Conducts hyperparameter tuning for the model, with an option to use Ray Tune.
+        Conduct hyperparameter tuning for the model, with an option to use Ray Tune.
 
         This method supports two modes of hyperparameter tuning: using Ray Tune or a custom tuning method.
         When Ray Tune is enabled, it leverages the 'run_ray_tune' function from the ultralytics.utils.tuner module.
@@ -881,7 +882,7 @@ class Model(torch.nn.Module):
     @property
     def names(self) -> Dict[int, str]:
         """
-        Retrieves the class names associated with the loaded model.
+        Retrieve the class names associated with the loaded model.
 
         This property returns the class names if they are defined in the model. It checks the class names for validity
         using the 'check_class_names' function from the ultralytics.nn.autobackend module. If the predictor is not
@@ -935,7 +936,7 @@ class Model(torch.nn.Module):
     @property
     def transforms(self):
         """
-        Retrieves the transformations applied to the input data of the loaded model.
+        Retrieve the transformations applied to the input data of the loaded model.
 
         This property returns the transformations if they are defined in the model. The transforms
         typically include preprocessing steps like resizing, normalization, and data augmentation
@@ -982,7 +983,7 @@ class Model(torch.nn.Module):
 
     def clear_callback(self, event: str) -> None:
         """
-        Clears all callback functions registered for a specified event.
+        Clear all callback functions registered for a specified event.
 
         This method removes all custom and default callback functions associated with the given event.
         It resets the callback list for the specified event to an empty list, effectively removing all
@@ -1062,7 +1063,7 @@ class Model(torch.nn.Module):
 
     def _smart_load(self, key: str):
         """
-        Intelligently loads the appropriate module based on the model task.
+        Intelligently load the appropriate module based on the model task.
 
         This method dynamically selects and returns the correct module (model, trainer, validator, or predictor)
         based on the current task of the model and the provided key. It uses the task_map dictionary to determine
@@ -1092,7 +1093,7 @@ class Model(torch.nn.Module):
     @property
     def task_map(self) -> dict:
         """
-        Provides a mapping from model tasks to corresponding classes for different modes.
+        Provide a mapping from model tasks to corresponding classes for different modes.
 
         This property method returns a dictionary that maps each supported task (e.g., detect, segment, classify)
         to a nested dictionary. The nested dictionary contains mappings for different operational modes