image-match/python/py/Lib/site-packages/timm/data/naflex_dataset.py

""" Dynamic Sequence Length Datasets for Variable Resolution Image Processing

Implements two dataset wrappers:
1. NaFlexMapDatasetWrapper - Map-style dataset that returns batches with variable sequence lengths
TODO: 2. NaFlexIterableDatasetWrapper - Iterable dataset that yields batches with variable sequence lengths

Both support:
- Pre-initialized transforms for efficiency
- Distributed training
- Multiple workers
- Variable batch sizes based on sequence length

Hacked together by / Copyright 2025, Ross Wightman, Hugging Face
"""

import math
import random
import warnings
from functools import partial
from typing import Any, Iterator, List, Tuple, Dict, Optional, Union, Callable

import torch
from torch.utils.data import Dataset, IterableDataset, DataLoader
from PIL import Image

from .naflex_transforms import Patchify
from timm.layers import to_2tuple


def calculate_naflex_batch_size(
        tokens_per_batch: int,
        seq_len: int,
        max_size: Optional[int] = None,
        divisor: int = 1,
        rounding: str = 'floor',
) -> int:
    """Calculate batch size based on sequence length with divisibility constraints.

    Args:
        tokens_per_batch: Target number of tokens per batch.
        seq_len: Sequence length for this batch.
        max_size: Optional maximum batch size.
        divisor: Ensure batch size is divisible by this value.
        rounding: Rounding method ('floor', 'ceil', 'round').

    Returns:
        Calculated batch size.
    """
    # Calculate raw batch size based on sequence length
    raw_batch_size = tokens_per_batch / seq_len

    # Apply divisibility with specified rounding method
    if divisor > 1:
        if rounding == 'floor':
            batch_size = math.floor(raw_batch_size / divisor) * divisor
        elif rounding == 'ceil':
            batch_size = math.ceil(raw_batch_size / divisor) * divisor
        else:  # 'round' is the default
            batch_size = round(raw_batch_size / divisor) * divisor
    else:
        # If no divisor specified, just use integer division
        batch_size = int(raw_batch_size)

    # Ensure batch size is valid
    batch_size = max(1, batch_size)  # At least 1

    if max_size is not None:
        batch_size = min(batch_size, max_size)

    return batch_size


class NaFlexCollator:
    """Custom collator for batching NaFlex-style variable-resolution images."""

    def __init__(
            self,
            max_seq_len: Optional[int] = None,
    ) -> None:
        """Initialize NaFlexCollator.

        Args:
            max_seq_len: Maximum sequence length for batching.
        """
        self.max_seq_len = max_seq_len or 576  # Default ViT-B/16 sequence length (577 = 24*24)

    def __call__(self, batch: List[Tuple[Dict[str, torch.Tensor], Union[int, torch.Tensor]]]) -> Tuple[Dict[str, torch.Tensor], torch.Tensor]:
        """Collate batch of NaFlex samples.

        Args:
            batch: List of tuples (patch_dict, target).

        Returns:
            A tuple of (input_dict, targets) where input_dict contains:
                - patches: Padded tensor of patches
                - patch_coord: Coordinates for each patch (y, x)
                - patch_valid: Valid indicators
        """
        assert isinstance(batch[0], tuple)
        batch_size = len(batch)

        # Extract targets
        targets = [item[1] for item in batch]
        if isinstance(targets[0], torch.Tensor):
            targets = torch.stack(targets)
        else:
            targets = torch.tensor(targets, dtype=torch.int64)

        # Get patch dictionaries
        patch_dicts = [item[0] for item in batch]

        # If we have a maximum sequence length constraint, ensure we don't exceed it
        if self.max_seq_len is not None:
            max_patches = self.max_seq_len
        else:
            # Find the maximum number of patches in this batch
            max_patches = max(item['patches'].shape[0] for item in patch_dicts)

        # Check if patches are flattened or unflattened
        patches_tensor = patch_dicts[0]['patches']
        is_unflattened = patches_tensor.ndim == 4  # [N, Ph, Pw, C]

        if is_unflattened:
            # Patches are [N, Ph, Pw, C] - variable patch size mode
            _, ph, pw, c = patches_tensor.shape
            patches = torch.zeros((batch_size, max_patches, ph, pw, c), dtype=torch.float32)
        else:
            # Patches are [N, P*P*C] - normal mode
            patch_dim = patches_tensor.shape[1]
            patches = torch.zeros((batch_size, max_patches, patch_dim), dtype=torch.float32)

        # Prepare other tensors
        patch_coord = torch.zeros((batch_size, max_patches, 2), dtype=torch.int64)  # [B, N, 2] for (y, x)
        patch_valid = torch.zeros((batch_size, max_patches), dtype=torch.bool)

        # Fill in the tensors
        for i, patch_dict in enumerate(patch_dicts):
            num_patches = min(patch_dict['patches'].shape[0], max_patches)

            patches[i, :num_patches] = patch_dict['patches'][:num_patches]
            patch_coord[i, :num_patches] = patch_dict['patch_coord'][:num_patches]
            patch_valid[i, :num_patches] = patch_dict['patch_valid'][:num_patches]

        result = {
            'patches': patches,
            'patch_coord': patch_coord,
            'patch_valid': patch_valid,
            'seq_len': max_patches,
        }

        return result, targets


def _resolve_patch_cfg(
        patch_size: Optional[Union[int, Tuple[int, int]]],
        patch_size_choices: Optional[List[int]],
        patch_size_choice_probs: Optional[List[float]],
) -> Tuple[List[Tuple[int, int]], List[float], bool]:
    """Resolve patch size configuration.

    Args:
        patch_size: Single patch size to use.
        patch_size_choices: List of patch sizes to choose from.
        patch_size_choice_probs: Probabilities for each patch size choice.

    Returns:
        Tuple of (sizes, probs, variable) where sizes is list of patch size tuples,
        probs is list of probabilities, and variable indicates if patch size varies.
    """
    # If both are None, default to patch_size=16
    if patch_size is None and patch_size_choices is None:
        patch_size = 16

    if (patch_size is None) == (patch_size_choices is None):
        raise ValueError(
            "Specify exactly one of `patch_size` or `patch_size_choices`."
        )

    if patch_size is not None:
        sizes = [to_2tuple(patch_size)]
        probs = [1.0]
        variable = False
    else:
        sizes = [to_2tuple(p) for p in patch_size_choices]
        if patch_size_choice_probs is None:
            probs = [1.0 / len(sizes)] * len(sizes)
        else:
            if len(patch_size_choice_probs) != len(sizes):
                raise ValueError("`patch_size_choice_probs` length mismatch.")
            s = float(sum(patch_size_choice_probs))
            if s <= 0:
                raise ValueError("`patch_size_choice_probs` sum to zero.")
            probs = [p / s for p in patch_size_choice_probs]
        variable = True
    return sizes, probs, variable


class NaFlexMapDatasetWrapper(IterableDataset):
    """
    IterableDataset wrapper for a map-style base dataset.

    Yields batches with variable sequence lengths. It calculates a canonical
    batch schedule (sequence length, batch size pairs) once based on the
    total dataset size (padded for distribution). Each epoch, it shuffles
    the order of this canonical schedule and the dataset indices.
    This ensures a consistent number of batches and samples per epoch
    across all ranks. Handles distributed training and multiple workers.
    """

    def __init__(
            self,
            base_dataset: Dataset,
            patch_size: Optional[Union[int, Tuple[int, int]]] = None,
            patch_size_choices: Optional[List[int]] = None,
            patch_size_choice_probs: Optional[List[float]] = None,
            seq_lens: Tuple[int, ...] = (128, 256, 576, 784, 1024),
            max_tokens_per_batch: int = 4096 * 4,
            transform_factory: Optional[Callable] = None,
            mixup_fn: Optional[Callable] = None,
            seed: int = 42,
            shuffle: bool = True,
            distributed: bool = False,
            rank: int = 0,
            world_size: int = 1,
            epoch: int = 0,
            batch_divisor: int = 8,
    ) -> None:
        """Initialize NaFlexMapDatasetWrapper.

        Args:
            base_dataset: Map-style dataset to wrap.
            patch_size: Single patch size to use.
            patch_size_choices: List of patch sizes to randomly select from.
            patch_size_choice_probs: Probabilities for each patch size.
            seq_lens: Sequence lengths to use for batching.
            max_tokens_per_batch: Target tokens per batch.
            transform_factory: Factory function for creating transforms.
            mixup_fn: Optional mixup function.
            seed: Random seed.
            shuffle: Whether to shuffle data.
            distributed: Whether using distributed training.
            rank: Process rank for distributed training.
            world_size: Total number of processes.
            epoch: Starting epoch.
            batch_divisor: Ensure batch size is divisible by this.
        """
        super().__init__()
        if not hasattr(base_dataset, '__len__') or not hasattr(base_dataset, '__getitem__'):
            raise TypeError("base_dataset must be a map-style dataset (implement __len__ and __getitem__)")

        self.base_dataset = base_dataset
        self.seq_lens = sorted(list(set(seq_lens))) # Ensure unique and sorted
        self.max_tokens_per_batch = max_tokens_per_batch
        self.seed = seed
        self.shuffle = shuffle
        self.distributed = distributed
        self.rank = rank if distributed else 0
        self.world_size = world_size if distributed else 1
        self.epoch = epoch
        self.batch_divisor = batch_divisor

        # Resolve patch size configuration
        self.patch_sizes, self.patch_size_probs, self.variable_patch_size = _resolve_patch_cfg(
            patch_size,
            patch_size_choices,
            patch_size_choice_probs
        )

        # Pre-initialize transforms and collate fns for each (seq_len, patch_idx) combination
        self.transforms: Dict[Tuple[int, int], Optional[Callable]] = {}
        self.collate_fns: Dict[int, Callable] = {}
        self.patchifiers: List[Callable] = []

        for seq_len in self.seq_lens:
            self.collate_fns[seq_len] = NaFlexCollator(seq_len)

        for patch_idx, patch_size_tuple in enumerate(self.patch_sizes):
            # Pre-initialize patchifiers for each patch size (indexed by patch_idx)
            self.patchifiers.append(Patchify(
                patch_size=patch_size_tuple,
                flatten_patches=not self.variable_patch_size
            ))

            # Create transforms for each (seq_len, patch_idx) combination
            for seq_len in self.seq_lens:
                key = (seq_len, patch_idx)
                if transform_factory:
                    self.transforms[key] = transform_factory(max_seq_len=seq_len, patch_size=patch_size_tuple)
                else:
                    self.transforms[key] = None # No transform

        self.mixup_fn = mixup_fn

        # Canonical Schedule Calculation (Done Once)
        self._canonical_batch_schedule: List[Tuple[int, int]] = []
        self._num_batches_per_rank: int = 0
        self._padded_samples_per_rank: int = 0
        self._create_canonical_schedule() # Calculate schedule based on padded size

        # Per-Epoch State
        # Stores (seq_len, list_of_indices) for the current epoch, specific to this rank
        self._epoch_batches: List[Tuple[int, List[int]]] = []
        self._prepare_epoch_batches(self.epoch)  # setup for initial epoch

    def _create_canonical_schedule(self):
        """
        Calculates the canonical batch schedule (seq_len, batch_size pairs)
        based on the dataset size, padded for distributed training.
        This schedule is the *same* for all ranks and ensures consistent
        epoch length. It is calculated once during initialization.
        """
        total_len = len(self.base_dataset)
        padded_total_len = total_len
        num_samples_per_rank = total_len

        if self.distributed and self.world_size > 1:
            # Calculate padding needed for even distribution
            if total_len % self.world_size != 0:
                 pad_size = self.world_size - (total_len % self.world_size)
                 padded_total_len += pad_size
                 print(f"Rank {self.rank}: Padding dataset with {pad_size} samples for distributed training (total size {padded_total_len}).")
            else:
                 pad_size = 0

            if padded_total_len % self.world_size != 0:
                 # This should not happen with the padding logic, but safeguard
                 raise RuntimeError(f"Internal Error: Padded total length {padded_total_len} not divisible by world size {self.world_size}")

            num_samples_per_rank = padded_total_len // self.world_size
        elif self.distributed and self.world_size <= 1:
             # Distributed flag set but world_size is 1, treat as non-distributed
             pass # num_samples_per_rank remains total_len

        self._padded_samples_per_rank = num_samples_per_rank

        if num_samples_per_rank == 0:
             self._canonical_batch_schedule = []
             self._num_batches_per_rank = 0
             return

        # Use a fixed seed for generating the canonical schedule structure
        g = torch.Generator()
        g.manual_seed(self.seed) # Use base seed, NOT epoch seed

        current_schedule: List[Tuple[int, int]] = []
        remaining_samples = num_samples_per_rank
        total_scheduled_samples = 0

        while remaining_samples > 0:
            # Sample sequence length deterministically based on base seed
            seq_idx = torch.randint(0, len(self.seq_lens), (1,), generator=g).item()
            seq_len = self.seq_lens[seq_idx]

            # Calculate batch size
            batch_size = calculate_naflex_batch_size(
                tokens_per_batch=self.max_tokens_per_batch,
                seq_len=seq_len,
                # max_size should be remaining_samples to avoid overshooting
                max_size=remaining_samples,
                divisor=self.batch_divisor,
                rounding='floor',
            )
            # Ensure batch size is positive and doesn't exceed remaining samples
            batch_size = max(1, batch_size)
            batch_size = min(batch_size, remaining_samples)

            if batch_size <= 0:
                 warnings.warn(f"Calculated batch size <= 0 (seq_len={seq_len}, remaining={remaining_samples}). Stopping schedule generation early.")
                 break # Avoid infinite loop if something goes wrong

            current_schedule.append((seq_len, batch_size))
            remaining_samples -= batch_size
            total_scheduled_samples += batch_size

        # Sanity check: Ensure the schedule covers all samples for the rank
        if total_scheduled_samples != num_samples_per_rank:
            warnings.warn(
                f"Rank {self.rank}: Canonical schedule accounts for {total_scheduled_samples} samples, "
                f"but expected {num_samples_per_rank} samples per rank. "
                f"This might happen if min_batch_size or batch_divisor constraints prevent utilizing all samples. "
                f"Check parameters. Remaining samples: {remaining_samples}"
            )
            # Adjust if needed? Could add a final small batch, but might violate constraints.
            # Current behavior: some samples might be dropped if schedule logic fails.

        self._canonical_batch_schedule = current_schedule
        self._num_batches_per_rank = len(current_schedule)
        print(f"Rank {self.rank}: Created canonical schedule with {self._num_batches_per_rank} batches for {self._padded_samples_per_rank} samples/rank.")


    def _prepare_epoch_batches(self, epoch: int):
        """
        Prepares the batches for the current epoch by:
        1. Shuffling the full dataset indices (using epoch seed).
        2. Applying padding if in distributed mode.
        3. Selecting indices for the current rank.
        4. Shuffling the *order* of the canonical batch schedule (using epoch seed).
        5. Assigning the rank's indices to the shuffled batches.
        """
        g = torch.Generator()
        g.manual_seed(self.seed + epoch) # Epoch-specific seed for shuffling

        # 1. Get shuffled global indices
        total_len = len(self.base_dataset)
        if self.shuffle:
            all_indices_shuffled = torch.randperm(total_len, generator=g).tolist()
        else:
            all_indices_shuffled = list(range(total_len))

        # 2. Apply padding for distributed mode
        indices_for_ranks = all_indices_shuffled
        if self.distributed and self.world_size > 1:
            padded_total_len = self._padded_samples_per_rank * self.world_size
            if padded_total_len > total_len:
                pad_size = padded_total_len - total_len
                # Repeat initial elements from the *shuffled* list for padding
                indices_for_ranks = all_indices_shuffled + all_indices_shuffled[:pad_size]
            # Ensure length matches expectation
            if len(indices_for_ranks) != padded_total_len:
                 raise RuntimeError(f"Internal Error: Padded index list length {len(indices_for_ranks)} does not match expected {padded_total_len}")

        # 3. Select indices for the current rank
        if self.distributed and self.world_size > 1:
            indices_this_rank = indices_for_ranks[self.rank::self.world_size]
        else: # Non-distributed or world_size=1
            indices_this_rank = indices_for_ranks

        # Sanity check length
        if len(indices_this_rank) != self._padded_samples_per_rank:
             # This might happen if canonical schedule generation had warnings/issues
             warnings.warn(
                 f"Rank {self.rank}: Number of indices for this rank ({len(indices_this_rank)}) "
                 f"does not match expected padded samples per rank ({self._padded_samples_per_rank}). "
                 f"Epoch generation might be inconsistent."
              )
             # Adjust expected samples? Or truncate/pad indices? Let's proceed but warn.
             # Using min() prevents IndexError later if indices are fewer than expected.
             effective_samples_this_rank = min(len(indices_this_rank), self._padded_samples_per_rank)
             indices_this_rank = indices_this_rank[:effective_samples_this_rank]

        else:
             effective_samples_this_rank = self._padded_samples_per_rank

        # 4. Shuffle the order of the canonical batch schedule for this epoch
        if self.shuffle:
            schedule_perm = torch.randperm(self._num_batches_per_rank, generator=g).tolist()
            shuffled_schedule = [self._canonical_batch_schedule[i] for i in schedule_perm]
        else:
            shuffled_schedule = list(self._canonical_batch_schedule) # Keep original order

        # 5. Assign indices to the shuffled batches
        self._epoch_batches = []
        idx_pos = 0
        scheduled_samples_count = 0
        for seq_len, bs in shuffled_schedule:
            # Ensure we don't try to grab more indices than available for the rank
            actual_bs = min(bs, effective_samples_this_rank - idx_pos)
            if actual_bs <= 0:
                 if scheduled_samples_count < effective_samples_this_rank:
                     # This indicates mismatch between schedule total and actual samples
                     warnings.warn(f"Rank {self.rank}: Ran out of samples ({idx_pos}/{effective_samples_this_rank}) before processing entire schedule. Check schedule generation.")
                 break # Stop if no more indices or batch size is zero

            batch_indices = indices_this_rank[idx_pos : idx_pos + actual_bs]
            self._epoch_batches.append((seq_len, batch_indices))
            idx_pos += actual_bs
            scheduled_samples_count += actual_bs

        # Final check
        if scheduled_samples_count != effective_samples_this_rank:
             warnings.warn(
                f"Rank {self.rank}: Assigned {scheduled_samples_count} samples to batches, "
                f"but expected {effective_samples_this_rank} effective samples this epoch. "
                f"Indices remaining: {effective_samples_this_rank - scheduled_samples_count}."
             )

    def set_epoch(self, epoch: int) -> None:
        """Updates the epoch, regenerating the epoch-specific batches.

        Args:
            epoch: New epoch number.
        """
        # Only regenerate if the epoch actually changes
        if epoch != self.epoch:
            self.epoch = epoch
            self._prepare_epoch_batches(epoch)

    def __len__(self) -> int:
        """Returns the number of batches per worker for the current epoch.

        Returns:
            Number of batches this worker will process.
        """
        return self._num_batches_per_rank

    def __iter__(self) -> Iterator[Tuple[Dict[str, torch.Tensor], torch.Tensor]]:
        """Iterates through pre-calculated batches for the current epoch.

        Yields:
            Tuple of (input_dict, targets) for each batch.
        """
        worker_info = torch.utils.data.get_worker_info()
        num_workers = worker_info.num_workers if worker_info else 1
        worker_id = worker_info.id if worker_info else 0

        # Distribute pre-calculated batches among workers for this rank
        # Each worker processes a slice of the batches prepared in _prepare_epoch_batches
        batches_for_worker = self._epoch_batches[worker_id::num_workers]
        for seq_len, indices in batches_for_worker:
            if not indices: # Skip if a batch ended up with no indices (shouldn't happen often)
                 continue

            # Select patch size for this batch
            patch_idx = 0
            if self.variable_patch_size:
                # Use torch multinomial for weighted random choice
                patch_idx = torch.multinomial(torch.tensor(self.patch_size_probs), 1).item()

            # Get the pre-initialized transform and patchifier using patch_idx
            transform_key = (seq_len, patch_idx)
            transform = self.transforms.get(transform_key)
            batch_patchifier = self.patchifiers[patch_idx]

            batch_imgs = []
            batch_targets = []
            for idx in indices:
                try:
                    # Get original image and label from map-style dataset
                    img, label = self.base_dataset[idx]

                    # Apply transform if available
                    # Handle cases where transform might return None or fail
                    processed_img = transform(img) if transform else img
                    if processed_img is None:
                        warnings.warn(f"Transform returned None for index {idx}. Skipping sample.")
                        continue

                    batch_imgs.append(processed_img)
                    batch_targets.append(label)

                except IndexError:
                     warnings.warn(f"IndexError encountered for index {idx} (possibly due to padding/repeated indices). Skipping sample.")
                     continue
                except Exception as e:
                    # Log other potential errors during data loading/processing
                    warnings.warn(f"Error processing sample index {idx}. Error: {e}. Skipping sample.")
                    continue # Skip problematic sample

            if self.mixup_fn is not None:
                batch_imgs, batch_targets = self.mixup_fn(batch_imgs, batch_targets)

            batch_imgs = [batch_patchifier(img) for img in batch_imgs]
            batch_samples = list(zip(batch_imgs, batch_targets))
            if batch_samples: # Only yield if we successfully processed samples
                # Collate the processed samples into a batch
                yield self.collate_fns[seq_len](batch_samples)

            # If batch_samples is empty after processing 'indices', an empty batch is skipped.