dgenerate-ultralytics-headless 8.3.194__py3-none-any.whl → 8.3.196__py3-none-any.whl

ultralytics/models/yolo/detect/train.py

@@ -8,16 +8,17 @@ from copy import copy
 from typing import Any
 
 import numpy as np
+import torch
 import torch.nn as nn
 
 from ultralytics.data import build_dataloader, build_yolo_dataset
 from ultralytics.engine.trainer import BaseTrainer
 from ultralytics.models import yolo
 from ultralytics.nn.tasks import DetectionModel
-from ultralytics.utils import LOGGER, RANK
+from ultralytics.utils import DEFAULT_CFG, LOGGER, RANK
 from ultralytics.utils.patches import override_configs
 from ultralytics.utils.plotting import plot_images, plot_labels, plot_results
-from ultralytics.utils.torch_utils import de_parallel, torch_distributed_zero_first
+from ultralytics.utils.torch_utils import torch_distributed_zero_first, unwrap_model
 
 
 class DetectionTrainer(BaseTrainer):
@@ -29,7 +30,7 @@ class DetectionTrainer(BaseTrainer):
 
     Attributes:
         model (DetectionModel): The YOLO detection model being trained.
-        data (Dict): Dictionary containing dataset information including class names and number of classes.
+        data (dict): Dictionary containing dataset information including class names and number of classes.
         loss_names (tuple): Names of the loss components used in training (box_loss, cls_loss, dfl_loss).
 
     Methods:
@@ -53,6 +54,18 @@ class DetectionTrainer(BaseTrainer):
         >>> trainer.train()
     """
 
+    def __init__(self, cfg=DEFAULT_CFG, overrides: dict[str, Any] | None = None, _callbacks=None):
+        """
+        Initialize a DetectionTrainer object for training YOLO object detection model training.
+
+        Args:
+            cfg (dict, optional): Default configuration dictionary containing training parameters.
+            overrides (dict, optional): Dictionary of parameter overrides for the default configuration.
+            _callbacks (list, optional): List of callback functions to be executed during training.
+        """
+        super().__init__(cfg, overrides, _callbacks)
+        self.dynamic_tensors = ["batch_idx", "cls", "bboxes"]
+
     def build_dataset(self, img_path: str, mode: str = "train", batch: int | None = None):
         """
         Build YOLO Dataset for training or validation.
@@ -65,7 +78,7 @@ class DetectionTrainer(BaseTrainer):
         Returns:
             (Dataset): YOLO dataset object configured for the specified mode.
         """
-        gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32)
+        gs = max(int(unwrap_model(self.model).stride.max() if self.model else 0), 32)
         return build_yolo_dataset(self.args, img_path, batch, self.data, mode=mode, rect=mode == "val", stride=gs)
 
     def get_dataloader(self, dataset_path: str, batch_size: int = 16, rank: int = 0, mode: str = "train"):
@@ -88,20 +101,29 @@ class DetectionTrainer(BaseTrainer):
         if getattr(dataset, "rect", False) and shuffle:
             LOGGER.warning("'rect=True' is incompatible with DataLoader shuffle, setting shuffle=False")
             shuffle = False
-        workers = self.args.workers if mode == "train" else self.args.workers * 2
-        return build_dataloader(dataset, batch_size, workers, shuffle, rank)  # return dataloader
+        return build_dataloader(
+            dataset,
+            batch=batch_size,
+            workers=self.args.workers if mode == "train" else self.args.workers * 2,
+            shuffle=shuffle,
+            rank=rank,
+            drop_last=self.args.compile and mode == "train",
+        )
 
     def preprocess_batch(self, batch: dict) -> dict:
         """
         Preprocess a batch of images by scaling and converting to float.
 
         Args:
-            batch (Dict): Dictionary containing batch data with 'img' tensor.
+            batch (dict): Dictionary containing batch data with 'img' tensor.
 
         Returns:
-            (Dict): Preprocessed batch with normalized images.
+            (dict): Preprocessed batch with normalized images.
         """
-        batch["img"] = batch["img"].to(self.device, non_blocking=True).float() / 255
+        for k, v in batch.items():
+            if isinstance(v, torch.Tensor):
+                batch[k] = v.to(self.device, non_blocking=True)
+        batch["img"] = batch["img"].float() / 255
         if self.args.multi_scale:
             imgs = batch["img"]
             sz = (
@@ -116,6 +138,10 @@ class DetectionTrainer(BaseTrainer):
                 ]  # new shape (stretched to gs-multiple)
                 imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False)
             batch["img"] = imgs
+
+        if self.args.compile:
+            for k in self.dynamic_tensors:
+                torch._dynamo.maybe_mark_dynamic(batch[k], 0)
         return batch
 
     def set_model_attributes(self):
@@ -158,11 +184,11 @@ class DetectionTrainer(BaseTrainer):
         Return a loss dict with labeled training loss items tensor.
 
         Args:
-            loss_items (List[float], optional): List of loss values.
+            loss_items (list[float], optional): List of loss values.
             prefix (str): Prefix for keys in the returned dictionary.
 
         Returns:
-            (Dict | List): Dictionary of labeled loss items if loss_items is provided, otherwise list of keys.
+            (dict | list): Dictionary of labeled loss items if loss_items is provided, otherwise list of keys.
         """
         keys = [f"{prefix}/{x}" for x in self.loss_names]
         if loss_items is not None:
@@ -186,7 +212,7 @@ class DetectionTrainer(BaseTrainer):
         Plot training samples with their annotations.
 
         Args:
-            batch (Dict[str, Any]): Dictionary containing batch data.
+            batch (dict[str, Any]): Dictionary containing batch data.
             ni (int): Number of iterations.
         """
         plot_images(

ultralytics/models/yolo/detect/val.py

@@ -27,13 +27,13 @@ class DetectionValidator(BaseValidator):
     Attributes:
         is_coco (bool): Whether the dataset is COCO.
         is_lvis (bool): Whether the dataset is LVIS.
-        class_map (List[int]): Mapping from model class indices to dataset class indices.
+        class_map (list[int]): Mapping from model class indices to dataset class indices.
         metrics (DetMetrics): Object detection metrics calculator.
         iouv (torch.Tensor): IoU thresholds for mAP calculation.
         niou (int): Number of IoU thresholds.
-        lb (List[Any]): List for storing ground truth labels for hybrid saving.
-        jdict (List[Dict[str, Any]]): List for storing JSON detection results.
-        stats (Dict[str, List[torch.Tensor]]): Dictionary for storing statistics during validation.
+        lb (list[Any]): List for storing ground truth labels for hybrid saving.
+        jdict (list[dict[str, Any]]): List for storing JSON detection results.
+        stats (dict[str, list[torch.Tensor]]): Dictionary for storing statistics during validation.
 
     Examples:
         >>> from ultralytics.models.yolo.detect import DetectionValidator
@@ -49,8 +49,8 @@ class DetectionValidator(BaseValidator):
         Args:
             dataloader (torch.utils.data.DataLoader, optional): Dataloader to use for validation.
             save_dir (Path, optional): Directory to save results.
-            args (Dict[str, Any], optional): Arguments for the validator.
-            _callbacks (List[Any], optional): List of callback functions.
+            args (dict[str, Any], optional): Arguments for the validator.
+            _callbacks (list[Any], optional): List of callback functions.
         """
         super().__init__(dataloader, save_dir, args, _callbacks)
         self.is_coco = False
@@ -66,16 +66,15 @@ class DetectionValidator(BaseValidator):
         Preprocess batch of images for YOLO validation.
 
         Args:
-            batch (Dict[str, Any]): Batch containing images and annotations.
+            batch (dict[str, Any]): Batch containing images and annotations.
 
         Returns:
-            (Dict[str, Any]): Preprocessed batch.
+            (dict[str, Any]): Preprocessed batch.
         """
-        batch["img"] = batch["img"].to(self.device, non_blocking=True)
+        for k, v in batch.items():
+            if isinstance(v, torch.Tensor):
+                batch[k] = v.to(self.device, non_blocking=True)
         batch["img"] = (batch["img"].half() if self.args.half else batch["img"].float()) / 255
-        for k in {"batch_idx", "cls", "bboxes"}:
-            batch[k] = batch[k].to(self.device, non_blocking=True)
-
         return batch
 
     def init_metrics(self, model: torch.nn.Module) -> None:
@@ -114,7 +113,7 @@ class DetectionValidator(BaseValidator):
             preds (torch.Tensor): Raw predictions from the model.
 
         Returns:
-            (List[Dict[str, torch.Tensor]]): Processed predictions after NMS, where each dict contains
+            (list[dict[str, torch.Tensor]]): Processed predictions after NMS, where each dict contains
                 'bboxes', 'conf', 'cls', and 'extra' tensors.
         """
         outputs = nms.non_max_suppression(
@@ -136,10 +135,10 @@ class DetectionValidator(BaseValidator):
 
         Args:
             si (int): Batch index.
-            batch (Dict[str, Any]): Batch data containing images and annotations.
+            batch (dict[str, Any]): Batch data containing images and annotations.
 
         Returns:
-            (Dict[str, Any]): Prepared batch with processed annotations.
+            (dict[str, Any]): Prepared batch with processed annotations.
         """
         idx = batch["batch_idx"] == si
         cls = batch["cls"][idx].squeeze(-1)
@@ -163,10 +162,10 @@ class DetectionValidator(BaseValidator):
         Prepare predictions for evaluation against ground truth.
 
         Args:
-            pred (Dict[str, torch.Tensor]): Post-processed predictions from the model.
+            pred (dict[str, torch.Tensor]): Post-processed predictions from the model.
 
         Returns:
-            (Dict[str, torch.Tensor]): Prepared predictions in native space.
+            (dict[str, torch.Tensor]): Prepared predictions in native space.
         """
         if self.args.single_cls:
             pred["cls"] *= 0
@@ -177,8 +176,8 @@ class DetectionValidator(BaseValidator):
         Update metrics with new predictions and ground truth.
 
         Args:
-            preds (List[Dict[str, torch.Tensor]]): List of predictions from the model.
-            batch (Dict[str, Any]): Batch data containing ground truth.
+            preds (list[dict[str, torch.Tensor]]): List of predictions from the model.
+            batch (dict[str, Any]): Batch data containing ground truth.
         """
         for si, pred in enumerate(preds):
             self.seen += 1
@@ -232,7 +231,7 @@ class DetectionValidator(BaseValidator):
         Calculate and return metrics statistics.
 
         Returns:
-            (Dict[str, Any]): Dictionary containing metrics results.
+            (dict[str, Any]): Dictionary containing metrics results.
         """
         self.metrics.process(save_dir=self.save_dir, plot=self.args.plots, on_plot=self.on_plot)
         self.metrics.clear_stats()
@@ -263,11 +262,11 @@ class DetectionValidator(BaseValidator):
         Return correct prediction matrix.
 
         Args:
-            preds (Dict[str, torch.Tensor]): Dictionary containing prediction data with 'bboxes' and 'cls' keys.
-            batch (Dict[str, Any]): Batch dictionary containing ground truth data with 'bboxes' and 'cls' keys.
+            preds (dict[str, torch.Tensor]): Dictionary containing prediction data with 'bboxes' and 'cls' keys.
+            batch (dict[str, Any]): Batch dictionary containing ground truth data with 'bboxes' and 'cls' keys.
 
         Returns:
-            (Dict[str, np.ndarray]): Dictionary containing 'tp' key with correct prediction matrix of shape (N, 10) for 10 IoU levels.
+            (dict[str, np.ndarray]): Dictionary containing 'tp' key with correct prediction matrix of shape (N, 10) for 10 IoU levels.
         """
         if len(batch["cls"]) == 0 or len(preds["cls"]) == 0:
             return {"tp": np.zeros((len(preds["cls"]), self.niou), dtype=bool)}
@@ -300,14 +299,16 @@ class DetectionValidator(BaseValidator):
             (torch.utils.data.DataLoader): Dataloader for validation.
         """
         dataset = self.build_dataset(dataset_path, batch=batch_size, mode="val")
-        return build_dataloader(dataset, batch_size, self.args.workers, shuffle=False, rank=-1)  # return dataloader
+        return build_dataloader(
+            dataset, batch_size, self.args.workers, shuffle=False, rank=-1, drop_last=self.args.compile
+        )
 
     def plot_val_samples(self, batch: dict[str, Any], ni: int) -> None:
         """
         Plot validation image samples.
 
         Args:
-            batch (Dict[str, Any]): Batch containing images and annotations.
+            batch (dict[str, Any]): Batch containing images and annotations.
             ni (int): Batch index.
         """
         plot_images(
@@ -325,8 +326,8 @@ class DetectionValidator(BaseValidator):
         Plot predicted bounding boxes on input images and save the result.
 
         Args:
-            batch (Dict[str, Any]): Batch containing images and annotations.
-            preds (List[Dict[str, torch.Tensor]]): List of predictions from the model.
+            batch (dict[str, Any]): Batch containing images and annotations.
+            preds (list[dict[str, torch.Tensor]]): List of predictions from the model.
             ni (int): Batch index.
             max_det (Optional[int]): Maximum number of detections to plot.
         """
@@ -352,9 +353,9 @@ class DetectionValidator(BaseValidator):
         Save YOLO detections to a txt file in normalized coordinates in a specific format.
 
         Args:
-            predn (Dict[str, torch.Tensor]): Dictionary containing predictions with keys 'bboxes', 'conf', and 'cls'.
+            predn (dict[str, torch.Tensor]): Dictionary containing predictions with keys 'bboxes', 'conf', and 'cls'.
             save_conf (bool): Whether to save confidence scores.
-            shape (Tuple[int, int]): Shape of the original image (height, width).
+            shape (tuple[int, int]): Shape of the original image (height, width).
             file (Path): File path to save the detections.
         """
         from ultralytics.engine.results import Results
@@ -371,9 +372,9 @@ class DetectionValidator(BaseValidator):
         Serialize YOLO predictions to COCO json format.
 
         Args:
-            predn (Dict[str, torch.Tensor]): Predictions dictionary containing 'bboxes', 'conf', and 'cls' keys
+            predn (dict[str, torch.Tensor]): Predictions dictionary containing 'bboxes', 'conf', and 'cls' keys
                 with bounding box coordinates, confidence scores, and class predictions.
-            pbatch (Dict[str, Any]): Batch dictionary containing 'imgsz', 'ori_shape', 'ratio_pad', and 'im_file'.
+            pbatch (dict[str, Any]): Batch dictionary containing 'imgsz', 'ori_shape', 'ratio_pad', and 'im_file'.
 
         Examples:
              >>> result = {
@@ -417,10 +418,10 @@ class DetectionValidator(BaseValidator):
         Evaluate YOLO output in JSON format and return performance statistics.
 
         Args:
-            stats (Dict[str, Any]): Current statistics dictionary.
+            stats (dict[str, Any]): Current statistics dictionary.
 
         Returns:
-            (Dict[str, Any]): Updated statistics dictionary with COCO/LVIS evaluation results.
+            (dict[str, Any]): Updated statistics dictionary with COCO/LVIS evaluation results.
         """
         pred_json = self.save_dir / "predictions.json"  # predictions
         anno_json = (
@@ -446,16 +447,16 @@ class DetectionValidator(BaseValidator):
         including mAP50, mAP50-95, and LVIS-specific metrics if applicable.
 
         Args:
-            stats (Dict[str, Any]): Dictionary to store computed metrics and statistics.
+            stats (dict[str, Any]): Dictionary to store computed metrics and statistics.
             pred_json (str | Path]): Path to JSON file containing predictions in COCO format.
             anno_json (str | Path]): Path to JSON file containing ground truth annotations in COCO format.
-            iou_types (str | List[str]]): IoU type(s) for evaluation. Can be single string or list of strings.
+            iou_types (str | list[str]]): IoU type(s) for evaluation. Can be single string or list of strings.
                 Common values include "bbox", "segm", "keypoints". Defaults to "bbox".
-            suffix (str | List[str]]): Suffix to append to metric names in stats dictionary. Should correspond
+            suffix (str | list[str]]): Suffix to append to metric names in stats dictionary. Should correspond
                 to iou_types if multiple types provided. Defaults to "Box".
 
         Returns:
-            (Dict[str, Any]): Updated stats dictionary containing the computed COCO/LVIS evaluation metrics.
+            (dict[str, Any]): Updated stats dictionary containing the computed COCO/LVIS evaluation metrics.
         """
         if self.args.save_json and (self.is_coco or self.is_lvis) and len(self.jdict):
             LOGGER.info(f"\nEvaluating faster-coco-eval mAP using {pred_json} and {anno_json}...")

ultralytics/models/yolo/model.py

@@ -185,7 +185,7 @@ class YOLOWorld(Model):
         Set the model's class names for detection.
 
         Args:
-            classes (List[str]): A list of categories i.e. ["person"].
+            classes (list[str]): A list of categories i.e. ["person"].
         """
         self.model.set_classes(classes)
         # Remove background if it's given
@@ -299,8 +299,8 @@ class YOLOE(Model):
         classification tasks. The model must be an instance of YOLOEModel.
 
         Args:
-            vocab (List[str]): Vocabulary list containing tokens or words used by the model for text processing.
-            names (List[str]): List of class names that the model can detect or classify.
+            vocab (list[str]): Vocabulary list containing tokens or words used by the model for text processing.
+            names (list[str]): List of class names that the model can detect or classify.
 
         Raises:
             AssertionError: If the model is not an instance of YOLOEModel.
@@ -322,7 +322,7 @@ class YOLOE(Model):
         Set the model's class names and embeddings for detection.
 
         Args:
-            classes (List[str]): A list of categories i.e. ["person"].
+            classes (list[str]): A list of categories i.e. ["person"].
             embeddings (torch.Tensor): Embeddings corresponding to the classes.
         """
         assert isinstance(self.model, YOLOEModel)
@@ -381,7 +381,7 @@ class YOLOE(Model):
                 directory paths, URL/YouTube streams, PIL images, numpy arrays, or webcam indices.
             stream (bool): Whether to stream the prediction results. If True, results are yielded as a
                 generator as they are computed.
-            visual_prompts (Dict[str, List]): Dictionary containing visual prompts for the model. Must include
+            visual_prompts (dict[str, list]): Dictionary containing visual prompts for the model. Must include
                 'bboxes' and 'cls' keys when non-empty.
             refer_image (str | PIL.Image | np.ndarray, optional): Reference image for visual prompts.
             predictor (callable, optional): Custom predictor function. If None, a predictor is automatically
@@ -389,7 +389,7 @@ class YOLOE(Model):
             **kwargs (Any): Additional keyword arguments passed to the predictor.
 
         Returns:
-            (List | generator): List of Results objects or generator of Results objects if stream=True.
+            (list | generator): List of Results objects or generator of Results objects if stream=True.
 
         Examples:
             >>> model = YOLOE("yoloe-11s-seg.pt")

ultralytics/models/yolo/obb/train.py

@@ -37,21 +37,12 @@ class OBBTrainer(yolo.detect.DetectionTrainer):
         """
         Initialize an OBBTrainer object for training Oriented Bounding Box (OBB) models.
 
-        This trainer extends the DetectionTrainer class to specialize in training models that detect oriented
-        bounding boxes. It automatically sets the task to 'obb' in the configuration.
-
         Args:
             cfg (dict, optional): Configuration dictionary for the trainer. Contains training parameters and
                 model configuration.
             overrides (dict, optional): Dictionary of parameter overrides for the configuration. Any values here
                 will take precedence over those in cfg.
-            _callbacks (List[Any], optional): List of callback functions to be invoked during training.
-
-        Examples:
-            >>> from ultralytics.models.yolo.obb import OBBTrainer
-            >>> args = dict(model="yolo11n-obb.pt", data="dota8.yaml", epochs=3)
-            >>> trainer = OBBTrainer(overrides=args)
-            >>> trainer.train()
+            _callbacks (list[Any], optional): List of callback functions to be invoked during training.
         """
         if overrides is None:
             overrides = {}

ultralytics/models/yolo/obb/val.py

@@ -77,13 +77,13 @@ class OBBValidator(DetectionValidator):
         Compute the correct prediction matrix for a batch of detections and ground truth bounding boxes.
 
         Args:
-            preds (Dict[str, torch.Tensor]): Prediction dictionary containing 'cls' and 'bboxes' keys with detected
+            preds (dict[str, torch.Tensor]): Prediction dictionary containing 'cls' and 'bboxes' keys with detected
                 class labels and bounding boxes.
-            batch (Dict[str, torch.Tensor]): Batch dictionary containing 'cls' and 'bboxes' keys with ground truth
+            batch (dict[str, torch.Tensor]): Batch dictionary containing 'cls' and 'bboxes' keys with ground truth
                 class labels and bounding boxes.
 
         Returns:
-            (Dict[str, np.ndarray]): Dictionary containing 'tp' key with the correct prediction matrix as a numpy
+            (dict[str, np.ndarray]): Dictionary containing 'tp' key with the correct prediction matrix as a numpy
                 array with shape (N, 10), which includes 10 IoU levels for each detection, indicating the accuracy
                 of predictions compared to the ground truth.
 
@@ -104,7 +104,7 @@ class OBBValidator(DetectionValidator):
             preds (torch.Tensor): Raw predictions from the model.
 
         Returns:
-            (List[Dict[str, torch.Tensor]]): Processed predictions with angle information concatenated to bboxes.
+            (list[dict[str, torch.Tensor]]): Processed predictions with angle information concatenated to bboxes.
         """
         preds = super().postprocess(preds)
         for pred in preds:
@@ -117,7 +117,7 @@ class OBBValidator(DetectionValidator):
 
         Args:
             si (int): Batch index to process.
-            batch (Dict[str, Any]): Dictionary containing batch data with keys:
+            batch (dict[str, Any]): Dictionary containing batch data with keys:
                 - batch_idx: Tensor of batch indices
                 - cls: Tensor of class labels
                 - bboxes: Tensor of bounding boxes
@@ -126,7 +126,7 @@ class OBBValidator(DetectionValidator):
                 - ratio_pad: Ratio and padding information
 
         Returns:
-            (Dict[str, Any]): Prepared batch data with scaled bounding boxes and metadata.
+            (dict[str, Any]): Prepared batch data with scaled bounding boxes and metadata.
         """
         idx = batch["batch_idx"] == si
         cls = batch["cls"][idx].squeeze(-1)
@@ -150,8 +150,8 @@ class OBBValidator(DetectionValidator):
         Plot predicted bounding boxes on input images and save the result.
 
         Args:
-            batch (Dict[str, Any]): Batch data containing images, file paths, and other metadata.
-            preds (List[torch.Tensor]): List of prediction tensors for each image in the batch.
+            batch (dict[str, Any]): Batch data containing images, file paths, and other metadata.
+            preds (list[torch.Tensor]): List of prediction tensors for each image in the batch.
             ni (int): Batch index used for naming the output file.
 
         Examples:
@@ -170,9 +170,9 @@ class OBBValidator(DetectionValidator):
         Convert YOLO predictions to COCO JSON format with rotated bounding box information.
 
         Args:
-            predn (Dict[str, torch.Tensor]): Prediction dictionary containing 'bboxes', 'conf', and 'cls' keys
+            predn (dict[str, torch.Tensor]): Prediction dictionary containing 'bboxes', 'conf', and 'cls' keys
                 with bounding box coordinates, confidence scores, and class predictions.
-            pbatch (Dict[str, Any]): Batch dictionary containing 'imgsz', 'ori_shape', 'ratio_pad', and 'im_file'.
+            pbatch (dict[str, Any]): Batch dictionary containing 'imgsz', 'ori_shape', 'ratio_pad', and 'im_file'.
 
         Notes:
             This method processes rotated bounding box predictions and converts them to both rbox format
@@ -204,7 +204,7 @@ class OBBValidator(DetectionValidator):
             predn (torch.Tensor): Predicted detections with shape (N, 7) containing bounding boxes, confidence scores,
                 class predictions, and angles in format (x, y, w, h, conf, cls, angle).
             save_conf (bool): Whether to save confidence scores in the text file.
-            shape (Tuple[int, int]): Original image shape in format (height, width).
+            shape (tuple[int, int]): Original image shape in format (height, width).
             file (Path): Output file path to save detections.
 
         Examples:
@@ -237,10 +237,10 @@ class OBBValidator(DetectionValidator):
         Evaluate YOLO output in JSON format and save predictions in DOTA format.
 
         Args:
-            stats (Dict[str, Any]): Performance statistics dictionary.
+            stats (dict[str, Any]): Performance statistics dictionary.
 
         Returns:
-            (Dict[str, Any]): Updated performance statistics.
+            (dict[str, Any]): Updated performance statistics.
         """
         if self.args.save_json and self.is_dota and len(self.jdict):
             import json

ultralytics/models/yolo/pose/train.py

@@ -44,9 +44,6 @@ class PoseTrainer(yolo.detect.DetectionTrainer):
         """
         Initialize a PoseTrainer object for training YOLO pose estimation models.
 
-        This initializes a trainer specialized for pose estimation tasks, setting the task to 'pose' and
-        handling specific configurations needed for keypoint detection models.
-
         Args:
             cfg (dict, optional): Default configuration dictionary containing training parameters.
             overrides (dict, optional): Dictionary of parameter overrides for the default configuration.
@@ -55,17 +52,12 @@ class PoseTrainer(yolo.detect.DetectionTrainer):
         Notes:
             This trainer will automatically set the task to 'pose' regardless of what is provided in overrides.
             A warning is issued when using Apple MPS device due to known bugs with pose models.
-
-        Examples:
-            >>> from ultralytics.models.yolo.pose import PoseTrainer
-            >>> args = dict(model="yolo11n-pose.pt", data="coco8-pose.yaml", epochs=3)
-            >>> trainer = PoseTrainer(overrides=args)
-            >>> trainer.train()
         """
         if overrides is None:
             overrides = {}
         overrides["task"] = "pose"
         super().__init__(cfg, overrides, _callbacks)
+        self.dynamic_tensors = ["batch_idx", "cls", "bboxes", "keypoints"]
 
         if isinstance(self.args.device, str) and self.args.device.lower() == "mps":
             LOGGER.warning(

ultralytics/models/yolo/pose/val.py

@@ -22,7 +22,7 @@ class PoseValidator(DetectionValidator):
 
     Attributes:
         sigma (np.ndarray): Sigma values for OKS calculation, either OKS_SIGMA or ones divided by number of keypoints.
-        kpt_shape (List[int]): Shape of the keypoints, typically [17, 3] for COCO format.
+        kpt_shape (list[int]): Shape of the keypoints, typically [17, 3] for COCO format.
         args (dict): Arguments for the validator including task set to "pose".
         metrics (PoseMetrics): Metrics object for pose evaluation.
 
@@ -86,7 +86,7 @@ class PoseValidator(DetectionValidator):
     def preprocess(self, batch: dict[str, Any]) -> dict[str, Any]:
         """Preprocess batch by converting keypoints data to float and moving it to the device."""
         batch = super().preprocess(batch)
-        batch["keypoints"] = batch["keypoints"].to(self.device, non_blocking=True).float()
+        batch["keypoints"] = batch["keypoints"].float()
         return batch
 
     def get_desc(self) -> str:
@@ -132,7 +132,7 @@ class PoseValidator(DetectionValidator):
                 bounding boxes, confidence scores, class predictions, and keypoint data.
 
         Returns:
-            (Dict[torch.Tensor]): Dict of processed prediction dictionaries, each containing:
+            (dict[torch.Tensor]): Dict of processed prediction dictionaries, each containing:
                 - 'bboxes': Bounding box coordinates
                 - 'conf': Confidence scores
                 - 'cls': Class predictions
@@ -154,10 +154,10 @@ class PoseValidator(DetectionValidator):
 
         Args:
             si (int): Batch index.
-            batch (Dict[str, Any]): Dictionary containing batch data with keys like 'keypoints', 'batch_idx', etc.
+            batch (dict[str, Any]): Dictionary containing batch data with keys like 'keypoints', 'batch_idx', etc.
 
         Returns:
-            (Dict[str, Any]): Prepared batch with keypoints scaled to original image dimensions.
+            (dict[str, Any]): Prepared batch with keypoints scaled to original image dimensions.
 
         Notes:
             This method extends the parent class's _prepare_batch method by adding keypoint processing.
@@ -177,13 +177,13 @@ class PoseValidator(DetectionValidator):
         Return correct prediction matrix by computing Intersection over Union (IoU) between detections and ground truth.
 
         Args:
-            preds (Dict[str, torch.Tensor]): Dictionary containing prediction data with keys 'cls' for class predictions
+            preds (dict[str, torch.Tensor]): Dictionary containing prediction data with keys 'cls' for class predictions
                 and 'keypoints' for keypoint predictions.
-            batch (Dict[str, Any]): Dictionary containing ground truth data with keys 'cls' for class labels,
+            batch (dict[str, Any]): Dictionary containing ground truth data with keys 'cls' for class labels,
                 'bboxes' for bounding boxes, and 'keypoints' for keypoint annotations.
 
         Returns:
-            (Dict[str, np.ndarray]): Dictionary containing the correct prediction matrix including 'tp_p' for pose
+            (dict[str, np.ndarray]): Dictionary containing the correct prediction matrix including 'tp_p' for pose
                 true positives across 10 IoU levels.
 
         Notes:
@@ -207,9 +207,9 @@ class PoseValidator(DetectionValidator):
         Save YOLO pose detections to a text file in normalized coordinates.
 
         Args:
-            predn (Dict[str, torch.Tensor]): Dictionary containing predictions with keys 'bboxes', 'conf', 'cls' and 'keypoints.
+            predn (dict[str, torch.Tensor]): Dictionary containing predictions with keys 'bboxes', 'conf', 'cls' and 'keypoints.
             save_conf (bool): Whether to save confidence scores.
-            shape (Tuple[int, int]): Shape of the original image (height, width).
+            shape (tuple[int, int]): Shape of the original image (height, width).
             file (Path): Output file path to save detections.
 
         Notes:
@@ -234,9 +234,9 @@ class PoseValidator(DetectionValidator):
         to COCO format, and appends the results to the internal JSON dictionary (self.jdict).
 
         Args:
-            predn (Dict[str, torch.Tensor]): Prediction dictionary containing 'bboxes', 'conf', 'cls',
+            predn (dict[str, torch.Tensor]): Prediction dictionary containing 'bboxes', 'conf', 'cls',
                 and 'keypoints' tensors.
-            pbatch (Dict[str, Any]): Batch dictionary containing 'imgsz', 'ori_shape', 'ratio_pad', and 'im_file'.
+            pbatch (dict[str, Any]): Batch dictionary containing 'imgsz', 'ori_shape', 'ratio_pad', and 'im_file'.
 
         Notes:
             The method extracts the image ID from the filename stem (either as an integer if numeric, or as a string),

ultralytics/models/yolo/segment/predict.py

@@ -71,13 +71,13 @@ class SegmentationPredictor(DetectionPredictor):
         Construct a list of result objects from the predictions.
 
         Args:
-            preds (List[torch.Tensor]): List of predicted bounding boxes, scores, and masks.
+            preds (list[torch.Tensor]): List of predicted bounding boxes, scores, and masks.
             img (torch.Tensor): The image after preprocessing.
-            orig_imgs (List[np.ndarray]): List of original images before preprocessing.
-            protos (List[torch.Tensor]): List of prototype masks.
+            orig_imgs (list[np.ndarray]): List of original images before preprocessing.
+            protos (list[torch.Tensor]): List of prototype masks.
 
         Returns:
-            (List[Results]): List of result objects containing the original images, image paths, class names,
+            (list[Results]): List of result objects containing the original images, image paths, class names,
                 bounding boxes, and masks.
         """
         return [

ultralytics/models/yolo/segment/train.py

@@ -19,7 +19,7 @@ class SegmentationTrainer(yolo.detect.DetectionTrainer):
     functionality including model initialization, validation, and visualization.
 
     Attributes:
-        loss_names (Tuple[str]): Names of the loss components used during training.
+        loss_names (tuple[str]): Names of the loss components used during training.
 
     Examples:
         >>> from ultralytics.models.yolo.segment import SegmentationTrainer
@@ -32,24 +32,16 @@ class SegmentationTrainer(yolo.detect.DetectionTrainer):
         """
         Initialize a SegmentationTrainer object.
 
-        This initializes a trainer for segmentation tasks, extending the detection trainer with segmentation-specific
-        functionality. It sets the task to 'segment' and prepares the trainer for training segmentation models.
-
         Args:
             cfg (dict): Configuration dictionary with default training settings.
             overrides (dict, optional): Dictionary of parameter overrides for the default configuration.
             _callbacks (list, optional): List of callback functions to be executed during training.
-
-        Examples:
-            >>> from ultralytics.models.yolo.segment import SegmentationTrainer
-            >>> args = dict(model="yolo11n-seg.pt", data="coco8-seg.yaml", epochs=3)
-            >>> trainer = SegmentationTrainer(overrides=args)
-            >>> trainer.train()
         """
         if overrides is None:
             overrides = {}
         overrides["task"] = "segment"
         super().__init__(cfg, overrides, _callbacks)
+        self.dynamic_tensors = ["batch_idx", "cls", "bboxes", "masks"]
 
     def get_model(self, cfg: dict | str | None = None, weights: str | Path | None = None, verbose: bool = True):
         """