dgenerate-ultralytics-headless 8.3.214__py3-none-any.whl → 8.4.7__py3-none-any.whl

ultralytics/data/base.py

@@ -21,11 +21,10 @@ from ultralytics.utils.patches import imread
 
 
 class BaseDataset(Dataset):
-    """
-    Base dataset class for loading and processing image data.
+    """Base dataset class for loading and processing image data.
 
-    This class provides core functionality for loading images, caching, and preparing data for training and inference
-    in object detection tasks.
+    This class provides core functionality for loading images, caching, and preparing data for training and inference in
+    object detection tasks.
 
     Attributes:
         img_path (str): Path to the folder containing images.
@@ -34,7 +33,8 @@ class BaseDataset(Dataset):
         single_cls (bool): Whether to treat all objects as a single class.
         prefix (str): Prefix to print in log messages.
         fraction (float): Fraction of dataset to utilize.
-        channels (int): Number of channels in the images (1 for grayscale, 3 for RGB).
+        channels (int): Number of channels in the images (1 for grayscale, 3 for color). Color images loaded with OpenCV
+            are in BGR channel order.
         cv2_flag (int): OpenCV flag for reading images.
         im_files (list[str]): List of image file paths.
         labels (list[dict]): List of label data dictionaries.
@@ -86,8 +86,7 @@ class BaseDataset(Dataset):
         fraction: float = 1.0,
         channels: int = 3,
     ):
-        """
-        Initialize BaseDataset with given configuration and options.
+        """Initialize BaseDataset with given configuration and options.
 
         Args:
             img_path (str | list[str]): Path to the folder containing images or list of image paths.
@@ -103,7 +102,8 @@ class BaseDataset(Dataset):
             single_cls (bool): If True, single class training is used.
             classes (list[int], optional): List of included classes.
             fraction (float): Fraction of dataset to utilize.
-            channels (int): Number of channels in the images (1 for grayscale, 3 for RGB).
+            channels (int): Number of channels in the images (1 for grayscale, 3 for color). Color images loaded with
+                OpenCV are in BGR channel order.
         """
         super().__init__()
         self.img_path = img_path
@@ -148,8 +148,7 @@ class BaseDataset(Dataset):
         self.transforms = self.build_transforms(hyp=hyp)
 
     def get_img_files(self, img_path: str | list[str]) -> list[str]:
-        """
-        Read image files from the specified path.
+        """Read image files from the specified path.
 
         Args:
             img_path (str | list[str]): Path or list of paths to image directories or files.
@@ -186,8 +185,7 @@ class BaseDataset(Dataset):
         return im_files
 
     def update_labels(self, include_class: list[int] | None) -> None:
-        """
-        Update labels to include only specified classes.
+        """Update labels to include only specified classes.
 
         Args:
             include_class (list[int], optional): List of classes to include. If None, all classes are included.
@@ -210,8 +208,7 @@ class BaseDataset(Dataset):
                 self.labels[i]["cls"][:, 0] = 0
 
     def load_image(self, i: int, rect_mode: bool = True) -> tuple[np.ndarray, tuple[int, int], tuple[int, int]]:
-        """
-        Load an image from dataset index 'i'.
+        """Load an image from dataset index 'i'.
 
         Args:
             i (int): Index of the image to load.
@@ -286,8 +283,7 @@ class BaseDataset(Dataset):
             np.save(f.as_posix(), imread(self.im_files[i]), allow_pickle=False)
 
     def check_cache_disk(self, safety_margin: float = 0.5) -> bool:
-        """
-        Check if there's enough disk space for caching images.
+        """Check if there's enough disk space for caching images.
 
         Args:
             safety_margin (float): Safety margin factor for disk space calculation.
@@ -307,10 +303,10 @@ class BaseDataset(Dataset):
             b += im.nbytes
             if not os.access(Path(im_file).parent, os.W_OK):
                 self.cache = None
-                LOGGER.warning(f"{self.prefix}Skipping caching images to disk, directory not writeable")
+                LOGGER.warning(f"{self.prefix}Skipping caching images to disk, directory not writable")
                 return False
         disk_required = b * self.ni / n * (1 + safety_margin)  # bytes required to cache dataset to disk
-        total, used, free = shutil.disk_usage(Path(self.im_files[0]).parent)
+        total, _used, free = shutil.disk_usage(Path(self.im_files[0]).parent)
         if disk_required > free:
             self.cache = None
             LOGGER.warning(
@@ -322,8 +318,7 @@ class BaseDataset(Dataset):
         return True
 
     def check_cache_ram(self, safety_margin: float = 0.5) -> bool:
-        """
-        Check if there's enough RAM for caching images.
+        """Check if there's enough RAM for caching images.
 
         Args:
             safety_margin (float): Safety margin factor for RAM calculation.
@@ -381,8 +376,7 @@ class BaseDataset(Dataset):
         return self.transforms(self.get_image_and_label(index))
 
     def get_image_and_label(self, index: int) -> dict[str, Any]:
-        """
-        Get and return label information from the dataset.
+        """Get and return label information from the dataset.
 
         Args:
             index (int): Index of the image to retrieve.
@@ -410,8 +404,7 @@ class BaseDataset(Dataset):
         return label
 
     def build_transforms(self, hyp: dict[str, Any] | None = None):
-        """
-        Users can customize augmentations here.
+        """Users can customize augmentations here.
 
         Examples:
             >>> if self.augment:
@@ -424,8 +417,7 @@ class BaseDataset(Dataset):
         raise NotImplementedError
 
     def get_labels(self) -> list[dict[str, Any]]:
-        """
-        Users can customize their own format here.
+        """Users can customize their own format here.
 
         Examples:
             Ensure output is a dictionary with the following keys:

ultralytics/data/build.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import math
 import os
 import random
 from collections.abc import Iterator
@@ -11,8 +12,9 @@ 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 torch.utils.data import Dataset, dataloader, distributed
 
 from ultralytics.cfg import IterableSimpleNamespace
 from ultralytics.data.dataset import GroundingDataset, YOLODataset, YOLOMultiModalDataset
@@ -33,8 +35,7 @@ from ultralytics.utils.torch_utils import TORCH_2_0
 
 
 class InfiniteDataLoader(dataloader.DataLoader):
-    """
-    Dataloader that reuses workers for infinite iteration.
+    """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 without recreating workers.
@@ -50,7 +51,7 @@ class InfiniteDataLoader(dataloader.DataLoader):
         reset: Reset the iterator, useful when modifying dataset settings during training.
 
     Examples:
-        Create an infinite dataloader for training
+        Create an infinite DataLoader for training
         >>> dataset = YOLODataset(...)
         >>> dataloader = InfiniteDataLoader(dataset, batch_size=16, shuffle=True)
         >>> for batch in dataloader:  # Infinite iteration
@@ -75,7 +76,7 @@ class InfiniteDataLoader(dataloader.DataLoader):
             yield next(self.iterator)
 
     def __del__(self):
-        """Ensure that workers are properly terminated when the dataloader is deleted."""
+        """Ensure that workers are properly terminated when the DataLoader is deleted."""
         try:
             if not hasattr(self.iterator, "_workers"):
                 return
@@ -92,11 +93,10 @@ class InfiniteDataLoader(dataloader.DataLoader):
 
 
 class _RepeatSampler:
-    """
-    Sampler that repeats forever for infinite iteration.
+    """Sampler that repeats forever for infinite iteration.
 
-    This sampler wraps another sampler and yields its contents indefinitely, allowing for infinite iteration
-    over a dataset without recreating the sampler.
+    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.
@@ -112,7 +112,109 @@ class _RepeatSampler:
             yield from iter(self.sampler)
 
 
-def seed_worker(worker_id: int):  # 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 (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: Dataset,
+        num_replicas: int | None = None,
+        batch_size: int | None = None,
+        rank: int | None = None,
+        shuffle: bool = False,
+    ) -> None:
+        """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.num_replicas = num_replicas
+        self.rank = rank
+        self.epoch = 0
+        self.shuffle = shuffle
+        self.total_size = len(dataset)
+        # ensure all ranks have a sample if batch size >= total size; degenerates to round-robin sampler
+        self.batch_size = 1 if batch_size >= self.total_size else batch_size
+        self.num_batches = math.ceil(self.total_size / self.batch_size)
+
+    def _get_rank_indices(self) -> tuple[int, int]:
+        """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) -> Iterator:
+        """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) -> int:
+        """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: int) -> None:
+        """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) -> None:
     """Set dataloader worker seed for reproducibility across worker processes."""
     worker_seed = torch.initial_seed() % 2**32
     np.random.seed(worker_seed)
@@ -128,7 +230,7 @@ def build_yolo_dataset(
     rect: bool = False,
     stride: int = 32,
     multi_modal: bool = False,
-):
+) -> Dataset:
     """Build and return a YOLO dataset based on configuration parameters."""
     dataset = YOLOMultiModalDataset if multi_modal else YOLODataset
     return dataset(
@@ -159,7 +261,7 @@ def build_grounding(
     rect: bool = False,
     stride: int = 32,
     max_samples: int = 80,
-):
+) -> Dataset:
     """Build and return a GroundingDataset based on configuration parameters."""
     return GroundingDataset(
         img_path=img_path,
@@ -181,9 +283,16 @@ def build_grounding(
     )
 
 
-def build_dataloader(dataset, batch: int, workers: int, shuffle: bool = True, rank: int = -1, drop_last: bool = False):
-    """
-    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,
+) -> InfiniteDataLoader:
+    """Create and return an InfiniteDataLoader or DataLoader for training or validation.
 
     Args:
         dataset (Dataset): Dataset to load data from.
@@ -192,6 +301,7 @@ def build_dataloader(dataset, batch: int, workers: int, shuffle: bool = True, ra
         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.
@@ -204,7 +314,13 @@ def build_dataloader(dataset, batch: int, workers: int, shuffle: bool = True, ra
     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(
@@ -214,7 +330,7 @@ def build_dataloader(dataset, batch: int, workers: int, shuffle: bool = True, ra
         num_workers=nw,
         sampler=sampler,
         prefetch_factor=4 if nw > 0 else None,  # increase over default 2
-        pin_memory=nd > 0,
+        pin_memory=nd > 0 and pin_memory,
         collate_fn=getattr(dataset, "collate_fn", None),
         worker_init_fn=seed_worker,
         generator=generator,
@@ -222,9 +338,10 @@ def build_dataloader(dataset, batch: int, workers: int, shuffle: bool = True, ra
     )
 
 
-def check_source(source):
-    """
-    Check the type of input source and return corresponding flag values.
+def check_source(
+    source: str | int | Path | list | tuple | np.ndarray | Image.Image | torch.Tensor,
+) -> tuple[Any, bool, bool, bool, bool, bool]:
+    """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.
@@ -271,12 +388,17 @@ def check_source(source):
     return source, webcam, screenshot, from_img, in_memory, tensor
 
 
-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.
+def load_inference_source(
+    source: str | int | Path | list | tuple | np.ndarray | Image.Image | torch.Tensor,
+    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.
+        source (str | Path | list | tuple | torch.Tensor | PIL.Image | np.ndarray): 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.
@@ -295,7 +417,7 @@ def load_inference_source(source=None, batch: int = 1, vid_stride: int = 1, buff
     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)
 
-    # Dataloader
+    # DataLoader
     if tensor:
         dataset = LoadTensor(source)
     elif in_memory: