timm/data/naflex_loader.py

"""NaFlex data loader for dynamic sequence length training.

This module provides a specialized data loader for Vision Transformer models that supports:
- Dynamic sequence length sampling during training for improved efficiency
- Variable patch size training with probabilistic selection
- Patch-level random erasing augmentation
- Efficient GPU prefetching with normalization

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

import math
from contextlib import suppress
from functools import partial
from typing import Callable, Dict, Iterator, List, Optional, Tuple, Union


import torch

from .constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from .loader import _worker_init, adapt_to_chs
from .naflex_dataset import NaFlexMapDatasetWrapper, NaFlexCollator
from .naflex_random_erasing import PatchRandomErasing
from .transforms_factory import create_transform


class NaFlexPrefetchLoader:
    """Data prefetcher for NaFlex format which normalizes patches."""

    def __init__(
            self,
            loader: torch.utils.data.DataLoader,
            mean: Tuple[float, ...] = IMAGENET_DEFAULT_MEAN,
            std: Tuple[float, ...] = IMAGENET_DEFAULT_STD,
            channels: int = 3,
            device: torch.device = torch.device('cuda'),
            img_dtype: Optional[torch.dtype] = None,
            re_prob: float = 0.,
            re_mode: str = 'const',
            re_count: int = 1,
            re_num_splits: int = 0,
        ) -> None:
        """Initialize NaFlexPrefetchLoader.

        Args:
            loader: DataLoader to prefetch from.
            mean: Mean values for normalization.
            std: Standard deviation values for normalization.
            channels: Number of image channels.
            device: Device to move tensors to.
            img_dtype: Data type for image tensors.
            re_prob: Random erasing probability.
            re_mode: Random erasing mode.
            re_count: Maximum number of erasing rectangles.
            re_num_splits: Number of augmentation splits.
        """
        self.loader = loader
        self.device = device
        self.img_dtype = img_dtype or torch.float32

        # Create mean/std tensors for normalization (will be applied to patches)
        mean = adapt_to_chs(mean, channels)
        std = adapt_to_chs(std, channels)
        normalization_shape = (1, 1, channels)
        self.channels = channels
        self.mean = torch.tensor(
            [x * 255 for x in mean], device=device, dtype=self.img_dtype).view(normalization_shape)
        self.std = torch.tensor(
            [x * 255 for x in std], device=device, dtype=self.img_dtype).view(normalization_shape)

        if re_prob > 0.:
            self.random_erasing = PatchRandomErasing(
                erase_prob=re_prob,
                mode=re_mode,
                max_count=re_count,
                num_splits=re_num_splits,
                device=device,
            )
        else:
            self.random_erasing = None

        # Check for CUDA/NPU availability
        self.is_cuda = device.type == 'cuda' and torch.cuda.is_available()
        self.is_npu = device.type == 'npu' and torch.npu.is_available()

    def __iter__(self) -> Iterator[Tuple[Dict[str, torch.Tensor], torch.Tensor]]:
        """Iterate through the loader with prefetching and normalization.

        Yields:
            Tuple of (input_dict, targets) with normalized patches.
        """
        first = True
        if self.is_cuda:
            stream = torch.cuda.Stream(device=self.device)
            stream_context = partial(torch.cuda.stream, stream=stream)
        elif self.is_npu:
            stream = torch.npu.Stream(device=self.device)
            stream_context = partial(torch.npu.stream, stream=stream)
        else:
            stream = None
            stream_context = suppress

        for next_input_dict, next_target in self.loader:
            with stream_context():
                # Move all tensors in input_dict to device
                for k, v in next_input_dict.items():
                    if isinstance(v, torch.Tensor):
                        dtype = self.img_dtype if k == 'patches' else None
                        next_input_dict[k] = next_input_dict[k].to(
                            device=self.device,
                            non_blocking=True,
                            dtype=dtype,
                        )

                next_target = next_target.to(device=self.device, non_blocking=True)

                # Normalize patch values - handle both [B, N, P*P*C] and [B, N, Ph, Pw, C] formats
                patches_tensor = next_input_dict['patches']
                original_shape = patches_tensor.shape

                if patches_tensor.ndim == 3:
                    # Format: [B, N, P*P*C] - flattened patches
                    batch_size, num_patches, patch_pixels = original_shape
                    # To [B*N, P*P, C] for normalization and erasing
                    patches = patches_tensor.view(batch_size, num_patches, -1, self.channels)
                elif patches_tensor.ndim == 5:
                    # Format: [B, N, Ph, Pw, C] - unflattened patches (variable patch size mode)
                    batch_size, num_patches, patch_h, patch_w, channels = original_shape
                    assert channels == self.channels, f"Expected {self.channels} channels, got {channels}"
                    # To [B*N, Ph*Pw, C] for normalization and erasing
                    patches = patches_tensor.view(batch_size, num_patches, -1, self.channels)
                else:
                    raise ValueError(f"Unexpected patches tensor dimensions: {patches_tensor.ndim}. Expected 3 or 5.")

                # Apply normalization
                patches = patches.sub(self.mean).div(self.std)

                if self.random_erasing is not None:
                    patches = self.random_erasing(
                        patches,
                        patch_coord=next_input_dict['patch_coord'],
                        patch_valid=next_input_dict.get('patch_valid', None),
                    )

                # Reshape back to original format
                next_input_dict['patches'] = patches.view(original_shape)

            if not first:
                yield input_dict, target
            else:
                first = False

            if stream is not None:
                if self.is_cuda:
                    torch.cuda.current_stream(device=self.device).wait_stream(stream)
                elif self.is_npu:
                    torch.npu.current_stream(device=self.device).wait_stream(stream)

            input_dict = next_input_dict
            target = next_target

        yield input_dict, target

    def __len__(self) -> int:
        """Get length of underlying loader.

        Returns:
            Number of batches in the loader.
        """
        return len(self.loader)

    @property
    def sampler(self):
        """Get sampler from underlying loader.

        Returns:
            Sampler from the underlying DataLoader.
        """
        return self.loader.sampler

    @property
    def dataset(self):
        """Get dataset from underlying loader.

        Returns:
            Dataset from the underlying DataLoader.
        """
        return self.loader.dataset


def create_naflex_loader(
        dataset,
        patch_size: Optional[Union[Tuple[int, int], int]] = None,
        patch_size_choices: Optional[List[int]] = None,
        patch_size_choice_probs: Optional[List[float]] = None,
        train_seq_lens: Tuple[int, ...] = (128, 256, 576, 784, 1024),
        max_seq_len: int = 576,
        batch_size: int = 32,
        is_training: bool = False,
        mixup_fn: Optional[Callable] = None,

        no_aug: bool = False,
        re_prob: float = 0.,
        re_mode: str = 'const',
        re_count: int = 1,
        re_split: bool = False,
        train_crop_mode: Optional[str] = None,
        scale: Optional[Tuple[float, float]] = None,
        ratio: Optional[Tuple[float, float]] = None,
        hflip: float = 0.5,
        vflip: float = 0.,
        color_jitter: float = 0.4,
        color_jitter_prob: Optional[float] = None,
        grayscale_prob: float = 0.,
        gaussian_blur_prob: float = 0.,
        auto_augment: Optional[str] = None,
        num_aug_repeats: int = 0,
        num_aug_splits: int = 0,
        interpolation: str = 'bilinear',
        mean: Tuple[float, ...] = IMAGENET_DEFAULT_MEAN,
        std: Tuple[float, ...] = IMAGENET_DEFAULT_STD,
        crop_pct: Optional[float] = None,
        crop_mode: Optional[str] = None,
        crop_border_pixels: Optional[int] = None,

        num_workers: int = 4,
        distributed: bool = False,
        rank: int = 0,
        world_size: int = 1,
        seed: int = 42,
        epoch: int = 0,
        use_prefetcher: bool = True,
        pin_memory: bool = True,
        img_dtype: torch.dtype = torch.float32,
        device: Union[str, torch.device] = torch.device('cuda'),
        persistent_workers: bool = True,
        worker_seeding: str = 'all',
    ) -> Union[torch.utils.data.DataLoader, NaFlexPrefetchLoader]:
    """Create a data loader with dynamic sequence length sampling for training.

    Args:
        dataset: Dataset to load from.
        patch_size: Single patch size to use.
        patch_size_choices: List of patch sizes for variable patch size training.
        patch_size_choice_probs: Probabilities for each patch size choice.
        train_seq_lens: Training sequence lengths for dynamic batching.
        max_seq_len: Fixed sequence length for validation.
        batch_size: Batch size for validation and max training sequence length.
        is_training: Whether this is for training (enables dynamic batching).
        mixup_fn: Optional mixup function.
        no_aug: Disable augmentation.
        re_prob: Random erasing probability.
        re_mode: Random erasing mode.
        re_count: Maximum number of erasing rectangles.
        re_split: Random erasing split flag.
        train_crop_mode: Training crop mode.
        scale: Scale range for random resize crop.
        ratio: Aspect ratio range for random resize crop.
        hflip: Horizontal flip probability.
        vflip: Vertical flip probability.
        color_jitter: Color jitter factor.
        color_jitter_prob: Color jitter probability.
        grayscale_prob: Grayscale conversion probability.
        gaussian_blur_prob: Gaussian blur probability.
        auto_augment: AutoAugment policy.
        num_aug_repeats: Number of augmentation repeats.
        num_aug_splits: Number of augmentation splits.
        interpolation: Interpolation method.
        mean: Normalization mean values.
        std: Normalization standard deviation values.
        crop_pct: Crop percentage for validation.
        crop_mode: Crop mode.
        crop_border_pixels: Crop border pixels.
        num_workers: Number of data loading workers.
        distributed: Whether using distributed training.
        rank: Process rank for distributed training.
        world_size: Total number of processes.
        seed: Random seed.
        epoch: Starting epoch.
        use_prefetcher: Whether to use prefetching.
        pin_memory: Whether to pin memory.
        img_dtype: Image data type.
        device: Device to move tensors to.
        persistent_workers: Whether to use persistent workers.
        worker_seeding: Worker seeding mode.

    Returns:
        DataLoader or NaFlexPrefetchLoader instance.
    """

    if is_training:
        # For training, use the dynamic sequence length mechanism
        assert num_aug_repeats == 0, 'Augmentation repeats not currently supported in NaFlex loader'

        transform_factory = partial(
            create_transform,
            is_training=True,
            no_aug=no_aug,
            train_crop_mode=train_crop_mode,
            scale=scale,
            ratio=ratio,
            hflip=hflip,
            vflip=vflip,
            color_jitter=color_jitter,
            color_jitter_prob=color_jitter_prob,
            grayscale_prob=grayscale_prob,
            gaussian_blur_prob=gaussian_blur_prob,
            auto_augment=auto_augment,
            interpolation=interpolation,
            mean=mean,
            std=std,
            crop_pct=crop_pct,
            crop_mode=crop_mode,
            crop_border_pixels=crop_border_pixels,
            re_prob=re_prob,
            re_mode=re_mode,
            re_count=re_count,
            use_prefetcher=use_prefetcher,
            naflex=True,
        )

        max_train_seq_len = max(train_seq_lens)
        max_tokens_per_batch = batch_size * max_train_seq_len

        if isinstance(dataset, torch.utils.data.IterableDataset):
            assert False, "IterableDataset Wrapper is a WIP"

        naflex_dataset = NaFlexMapDatasetWrapper(
            dataset,
            transform_factory=transform_factory,
            patch_size=patch_size,
            patch_size_choices=patch_size_choices,
            patch_size_choice_probs=patch_size_choice_probs,
            seq_lens=train_seq_lens,
            max_tokens_per_batch=max_tokens_per_batch,
            mixup_fn=mixup_fn,
            seed=seed,
            distributed=distributed,
            rank=rank,
            world_size=world_size,
            shuffle=True,
            epoch=epoch,
        )

        # NOTE: Collation is handled by the dataset wrapper for training
        loader = torch.utils.data.DataLoader(
            naflex_dataset,
            batch_size=None,
            shuffle=False,
            num_workers=num_workers,
            sampler=None,
            pin_memory=pin_memory,
            worker_init_fn=partial(_worker_init, worker_seeding=worker_seeding),
            persistent_workers=persistent_workers
        )

        if use_prefetcher:
            loader = NaFlexPrefetchLoader(
                loader,
                mean=mean,
                std=std,
                img_dtype=img_dtype,
                device=device,
                re_prob=re_prob,
                re_mode=re_mode,
                re_count=re_count,
            )

    else:
        # For validation, use fixed sequence length (unchanged)
        dataset.transform = create_transform(
            is_training=False,
            interpolation=interpolation,
            mean=mean,
            std=std,
            # FIXME add crop args when sequence transforms support crop modes
            use_prefetcher=use_prefetcher,
            naflex=True,
            patch_size=patch_size,
            max_seq_len=max_seq_len,
            patchify=True,
        )

        # Create the collator
        collate_fn = NaFlexCollator(max_seq_len=max_seq_len)

        # Handle distributed training
        sampler = None
        if distributed and not isinstance(dataset, torch.utils.data.IterableDataset):
            # For validation, use OrderedDistributedSampler
            from timm.data.distributed_sampler import OrderedDistributedSampler
            sampler = OrderedDistributedSampler(dataset)

        loader = torch.utils.data.DataLoader(
            dataset,
            batch_size=batch_size,
            shuffle=False,
            num_workers=num_workers,
            sampler=sampler,
            collate_fn=collate_fn,
            pin_memory=pin_memory,
            drop_last=False,
        )

        if use_prefetcher:
            loader = NaFlexPrefetchLoader(
                loader,
                mean=mean,
                std=std,
                img_dtype=img_dtype,
                device=device,
            )

    return loader