dgenerate-ultralytics-headless 8.3.137__py3-none-any.whl → 8.3.224__py3-none-any.whl

ultralytics/data/build.py

@@ -1,14 +1,22 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
+import math
 import os
 import random
+from collections.abc import Iterator
 from pathlib import Path
+from typing import Any
+from urllib.parse import urlsplit
 
 import numpy as np
 import torch
+import torch.distributed as dist
 from PIL import Image
 from torch.utils.data import dataloader, distributed
 
+from ultralytics.cfg import IterableSimpleNamespace
 from ultralytics.data.dataset import GroundingDataset, YOLODataset, YOLOMultiModalDataset
 from ultralytics.data.loaders import (
     LOADERS,
@@ -20,40 +28,49 @@ from ultralytics.data.loaders import (
     SourceTypes,
     autocast_list,
 )
-from ultralytics.data.utils import IMG_FORMATS, PIN_MEMORY, VID_FORMATS
+from ultralytics.data.utils import IMG_FORMATS, VID_FORMATS
 from ultralytics.utils import RANK, colorstr
 from ultralytics.utils.checks import check_file
+from ultralytics.utils.torch_utils import TORCH_2_0
 
 
 class InfiniteDataLoader(dataloader.DataLoader):
-    """
-    Dataloader that reuses workers.
+    """Dataloader that reuses workers for infinite iteration.
 
     This dataloader extends the PyTorch DataLoader to provide infinite recycling of workers, which improves efficiency
-    for training loops that need to iterate through the dataset multiple times.
+    for training loops that need to iterate through the dataset multiple times without recreating workers.
 
     Attributes:
         batch_sampler (_RepeatSampler): A sampler that repeats indefinitely.
         iterator (Iterator): The iterator from the parent DataLoader.
 
     Methods:
-        __len__: Returns the length of the batch sampler's sampler.
-        __iter__: Creates a sampler that repeats indefinitely.
-        __del__: Ensures workers are properly terminated.
-        reset: Resets the iterator, useful when modifying dataset settings during training.
+        __len__: Return the length of the batch sampler's sampler.
+        __iter__: Create a sampler that repeats indefinitely.
+        __del__: Ensure workers are properly terminated.
+        reset: Reset the iterator, useful when modifying dataset settings during training.
+
+    Examples:
+        Create an infinite dataloader for training
+        >>> dataset = YOLODataset(...)
+        >>> dataloader = InfiniteDataLoader(dataset, batch_size=16, shuffle=True)
+        >>> for batch in dataloader:  # Infinite iteration
+        >>>     train_step(batch)
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any):
         """Initialize the InfiniteDataLoader with the same arguments as DataLoader."""
+        if not TORCH_2_0:
+            kwargs.pop("prefetch_factor", None)  # not supported by earlier versions
         super().__init__(*args, **kwargs)
         object.__setattr__(self, "batch_sampler", _RepeatSampler(self.batch_sampler))
         self.iterator = super().__iter__()
 
-    def __len__(self):
+    def __len__(self) -> int:
         """Return the length of the batch sampler's sampler."""
         return len(self.batch_sampler.sampler)
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator:
         """Create an iterator that yields indefinitely from the underlying iterator."""
         for _ in range(len(self)):
             yield next(self.iterator)
@@ -76,34 +93,137 @@ class InfiniteDataLoader(dataloader.DataLoader):
 
 
 class _RepeatSampler:
-    """
-    Sampler that repeats forever.
+    """Sampler that repeats forever for infinite iteration.
 
-    This sampler wraps another sampler and yields its contents indefinitely, allowing for infinite iteration
-    over a dataset.
+    This sampler wraps another sampler and yields its contents indefinitely, allowing for infinite iteration over a
+    dataset without recreating the sampler.
 
     Attributes:
         sampler (Dataset.sampler): The sampler to repeat.
     """
 
-    def __init__(self, sampler):
+    def __init__(self, sampler: Any):
         """Initialize the _RepeatSampler with a sampler to repeat indefinitely."""
         self.sampler = sampler
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator:
         """Iterate over the sampler indefinitely, yielding its contents."""
         while True:
             yield from iter(self.sampler)
 
 
-def seed_worker(worker_id):  # noqa
+class ContiguousDistributedSampler(torch.utils.data.Sampler):
+    """Distributed sampler that assigns contiguous batch-aligned chunks of the dataset to each GPU.
+
+    Unlike PyTorch's DistributedSampler which distributes samples in a round-robin fashion (GPU 0 gets indices
+    [0,2,4,...], GPU 1 gets [1,3,5,...]), this sampler gives each GPU contiguous batches of the dataset (GPU 0 gets
+    batches [0,1,2,...], GPU 1 gets batches [k,k+1,...], etc.). This preserves any ordering or grouping in the original
+    dataset, which is critical when samples are organized by similarity (e.g., images sorted by size to enable efficient
+    batching without padding when using rect=True).
+
+    The sampler handles uneven batch counts by distributing remainder batches to the first few ranks, ensuring all
+    samples are covered exactly once across all GPUs.
+
+    Args:
+        dataset (torch.utils.data.Dataset): Dataset to sample from. Must implement __len__.
+        num_replicas (int, optional): Number of distributed processes. Defaults to world size.
+        batch_size (int, optional): Batch size used by dataloader. Defaults to dataset batch size.
+        rank (int, optional): Rank of current process. Defaults to current rank.
+        shuffle (bool, optional): Whether to shuffle indices within each rank's chunk. Defaults to False. When True,
+            shuffling is deterministic and controlled by set_epoch() for reproducibility.
+
+    Examples:
+        >>> # For validation with size-grouped images
+        >>> sampler = ContiguousDistributedSampler(val_dataset, batch_size=32, shuffle=False)
+        >>> loader = DataLoader(val_dataset, batch_size=32, sampler=sampler)
+        >>> # For training with shuffling
+        >>> sampler = ContiguousDistributedSampler(train_dataset, batch_size=32, shuffle=True)
+        >>> for epoch in range(num_epochs):
+        ...     sampler.set_epoch(epoch)
+        ...     for batch in loader:
+        ...         ...
+    """
+
+    def __init__(self, dataset, num_replicas=None, batch_size=None, rank=None, shuffle=False):
+        """Initialize the sampler with dataset and distributed training parameters."""
+        if num_replicas is None:
+            num_replicas = dist.get_world_size() if dist.is_initialized() else 1
+        if rank is None:
+            rank = dist.get_rank() if dist.is_initialized() else 0
+        if batch_size is None:
+            batch_size = getattr(dataset, "batch_size", 1)
+
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.batch_size = batch_size
+        self.rank = rank
+        self.epoch = 0
+        self.shuffle = shuffle
+        self.total_size = len(dataset)
+        self.num_batches = math.ceil(self.total_size / self.batch_size)
+
+    def _get_rank_indices(self):
+        """Calculate the start and end sample indices for this rank."""
+        # Calculate which batches this rank handles
+        batches_per_rank_base = self.num_batches // self.num_replicas
+        remainder = self.num_batches % self.num_replicas
+
+        # This rank gets an extra batch if rank < remainder
+        batches_for_this_rank = batches_per_rank_base + (1 if self.rank < remainder else 0)
+
+        # Calculate starting batch: base position + number of extra batches given to earlier ranks
+        start_batch = self.rank * batches_per_rank_base + min(self.rank, remainder)
+        end_batch = start_batch + batches_for_this_rank
+
+        # Convert batch indices to sample indices
+        start_idx = start_batch * self.batch_size
+        end_idx = min(end_batch * self.batch_size, self.total_size)
+
+        return start_idx, end_idx
+
+    def __iter__(self):
+        """Generate indices for this rank's contiguous chunk of the dataset."""
+        start_idx, end_idx = self._get_rank_indices()
+        indices = list(range(start_idx, end_idx))
+
+        if self.shuffle:
+            g = torch.Generator()
+            g.manual_seed(self.epoch)
+            indices = [indices[i] for i in torch.randperm(len(indices), generator=g).tolist()]
+
+        return iter(indices)
+
+    def __len__(self):
+        """Return the number of samples in this rank's chunk."""
+        start_idx, end_idx = self._get_rank_indices()
+        return end_idx - start_idx
+
+    def set_epoch(self, epoch):
+        """Set the epoch for this sampler to ensure different shuffling patterns across epochs.
+
+        Args:
+            epoch (int): Epoch number to use as the random seed for shuffling.
+        """
+        self.epoch = epoch
+
+
+def seed_worker(worker_id: int):
     """Set dataloader worker seed for reproducibility across worker processes."""
     worker_seed = torch.initial_seed() % 2**32
     np.random.seed(worker_seed)
     random.seed(worker_seed)
 
 
-def build_yolo_dataset(cfg, img_path, batch, data, mode="train", rect=False, stride=32, multi_modal=False):
+def build_yolo_dataset(
+    cfg: IterableSimpleNamespace,
+    img_path: str,
+    batch: int,
+    data: dict[str, Any],
+    mode: str = "train",
+    rect: bool = False,
+    stride: int = 32,
+    multi_modal: bool = False,
+):
     """Build and return a YOLO dataset based on configuration parameters."""
     dataset = YOLOMultiModalDataset if multi_modal else YOLODataset
     return dataset(
@@ -115,7 +235,7 @@ def build_yolo_dataset(cfg, img_path, batch, data, mode="train", rect=False, str
         rect=cfg.rect or rect,  # rectangular batches
         cache=cfg.cache or None,
         single_cls=cfg.single_cls or False,
-        stride=int(stride),
+        stride=stride,
         pad=0.0 if mode == "train" else 0.5,
         prefix=colorstr(f"{mode}: "),
         task=cfg.task,
@@ -125,11 +245,21 @@ def build_yolo_dataset(cfg, img_path, batch, data, mode="train", rect=False, str
     )
 
 
-def build_grounding(cfg, img_path, json_file, batch, mode="train", rect=False, stride=32):
+def build_grounding(
+    cfg: IterableSimpleNamespace,
+    img_path: str,
+    json_file: str,
+    batch: int,
+    mode: str = "train",
+    rect: bool = False,
+    stride: int = 32,
+    max_samples: int = 80,
+):
     """Build and return a GroundingDataset based on configuration parameters."""
     return GroundingDataset(
         img_path=img_path,
         json_file=json_file,
+        max_samples=max_samples,
         imgsz=cfg.imgsz,
         batch_size=batch,
         augment=mode == "train",  # augmentation
@@ -137,7 +267,7 @@ def build_grounding(cfg, img_path, json_file, batch, mode="train", rect=False, s
         rect=cfg.rect or rect,  # rectangular batches
         cache=cfg.cache or None,
         single_cls=cfg.single_cls or False,
-        stride=int(stride),
+        stride=stride,
         pad=0.0 if mode == "train" else 0.5,
         prefix=colorstr(f"{mode}: "),
         task=cfg.task,
@@ -146,24 +276,44 @@ def build_grounding(cfg, img_path, json_file, batch, mode="train", rect=False, s
     )
 
 
-def build_dataloader(dataset, batch, workers, shuffle=True, rank=-1):
-    """
-    Create and return an InfiniteDataLoader or DataLoader for training or validation.
+def build_dataloader(
+    dataset,
+    batch: int,
+    workers: int,
+    shuffle: bool = True,
+    rank: int = -1,
+    drop_last: bool = False,
+    pin_memory: bool = True,
+):
+    """Create and return an InfiniteDataLoader or DataLoader for training or validation.
 
     Args:
         dataset (Dataset): Dataset to load data from.
         batch (int): Batch size for the dataloader.
         workers (int): Number of worker threads for loading data.
-        shuffle (bool): Whether to shuffle the dataset.
-        rank (int): Process rank in distributed training. -1 for single-GPU training.
+        shuffle (bool, optional): Whether to shuffle the dataset.
+        rank (int, optional): Process rank in distributed training. -1 for single-GPU training.
+        drop_last (bool, optional): Whether to drop the last incomplete batch.
+        pin_memory (bool, optional): Whether to use pinned memory for dataloader.
 
     Returns:
         (InfiniteDataLoader): A dataloader that can be used for training or validation.
+
+    Examples:
+        Create a dataloader for training
+        >>> dataset = YOLODataset(...)
+        >>> dataloader = build_dataloader(dataset, batch=16, workers=4, shuffle=True)
     """
     batch = min(batch, len(dataset))
     nd = torch.cuda.device_count()  # number of CUDA devices
     nw = min(os.cpu_count() // max(nd, 1), workers)  # number of workers
-    sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle)
+    sampler = (
+        None
+        if rank == -1
+        else distributed.DistributedSampler(dataset, shuffle=shuffle)
+        if shuffle
+        else ContiguousDistributedSampler(dataset)
+    )
     generator = torch.Generator()
     generator.manual_seed(6148914691236517205 + RANK)
     return InfiniteDataLoader(
@@ -172,38 +322,46 @@ def build_dataloader(dataset, batch, workers, shuffle=True, rank=-1):
         shuffle=shuffle and sampler is None,
         num_workers=nw,
         sampler=sampler,
-        pin_memory=PIN_MEMORY,
+        prefetch_factor=4 if nw > 0 else None,  # increase over default 2
+        pin_memory=nd > 0 and pin_memory,
         collate_fn=getattr(dataset, "collate_fn", None),
         worker_init_fn=seed_worker,
         generator=generator,
+        drop_last=drop_last and len(dataset) % batch != 0,
     )
 
 
 def check_source(source):
-    """
-    Check the type of input source and return corresponding flag values.
+    """Check the type of input source and return corresponding flag values.
 
     Args:
-        source (str | int | Path | List | Tuple | np.ndarray | PIL.Image | torch.Tensor): The input source to check.
+        source (str | int | Path | list | tuple | np.ndarray | PIL.Image | torch.Tensor): The input source to check.
 
     Returns:
-        source (str | int | Path | List | Tuple | np.ndarray | PIL.Image | torch.Tensor): The processed source.
+        source (str | int | Path | list | tuple | np.ndarray | PIL.Image | torch.Tensor): The processed source.
         webcam (bool): Whether the source is a webcam.
         screenshot (bool): Whether the source is a screenshot.
         from_img (bool): Whether the source is an image or list of images.
         in_memory (bool): Whether the source is an in-memory object.
         tensor (bool): Whether the source is a torch.Tensor.
 
-    Raises:
-        TypeError: If the source type is unsupported.
+    Examples:
+        Check a file path source
+        >>> source, webcam, screenshot, from_img, in_memory, tensor = check_source("image.jpg")
+
+        Check a webcam source
+        >>> source, webcam, screenshot, from_img, in_memory, tensor = check_source(0)
     """
     webcam, screenshot, from_img, in_memory, tensor = False, False, False, False, False
     if isinstance(source, (str, int, Path)):  # int for local usb camera
         source = str(source)
-        is_file = Path(source).suffix[1:] in (IMG_FORMATS | VID_FORMATS)
-        is_url = source.lower().startswith(("https://", "http://", "rtsp://", "rtmp://", "tcp://"))
+        source_lower = source.lower()
+        is_url = source_lower.startswith(("https://", "http://", "rtsp://", "rtmp://", "tcp://"))
+        is_file = (urlsplit(source_lower).path if is_url else source_lower).rpartition(".")[-1] in (
+            IMG_FORMATS | VID_FORMATS
+        )
         webcam = source.isnumeric() or source.endswith(".streams") or (is_url and not is_file)
-        screenshot = source.lower() == "screen"
+        screenshot = source_lower == "screen"
         if is_url and is_file:
             source = check_file(source)  # download
     elif isinstance(source, LOADERS):
@@ -221,19 +379,25 @@ def check_source(source):
     return source, webcam, screenshot, from_img, in_memory, tensor
 
 
-def load_inference_source(source=None, batch=1, vid_stride=1, buffer=False, channels=3):
-    """
-    Load an inference source for object detection and apply necessary transformations.
+def load_inference_source(source=None, batch: int = 1, vid_stride: int = 1, buffer: bool = False, channels: int = 3):
+    """Load an inference source for object detection and apply necessary transformations.
 
     Args:
         source (str | Path | torch.Tensor | PIL.Image | np.ndarray, optional): The input source for inference.
         batch (int, optional): Batch size for dataloaders.
         vid_stride (int, optional): The frame interval for video sources.
         buffer (bool, optional): Whether stream frames will be buffered.
-        channels (int): The number of input channels for the model.
+        channels (int, optional): The number of input channels for the model.
 
     Returns:
         (Dataset): A dataset object for the specified input source with attached source_type attribute.
+
+    Examples:
+        Load an image source for inference
+        >>> dataset = load_inference_source("image.jpg", batch=1)
+
+        Load a video stream source
+        >>> dataset = load_inference_source("rtsp://example.com/stream", vid_stride=2)
     """
     source, stream, screenshot, from_img, in_memory, tensor = check_source(source)
     source_type = source.source_type if in_memory else SourceTypes(stream, screenshot, from_img, tensor)