dataloader.py

import math
import os

import torch
import torch.utils.data as data
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data.distributed import DistributedSampler

from dataset import COCODataset


def collate_fn(data_batch, pad_token_id=0):
    """Pads variable-length captions and stacks image tensors.

    Args:
        data_batch: List of tuples (image, caption[, image_id]) from the dataset.
        pad_token_id: The token ID to use for padding (from tokenizer.pad_token_id).

    Returns:
        images: Batched image tensors.
        captions: Padded caption tensors.
        lengths: Original caption lengths before padding.
        image_ids (optional): Tensor of image IDs when provided by the dataset.
    """
    # Sort batch by caption length (descending) for efficient RNN processing
    data_batch.sort(key=lambda x: len(x[1]), reverse=True)

    include_ids = len(data_batch[0]) == 3
    if include_ids:
        images, captions, image_ids = zip(*data_batch)
    else:
        images, captions = zip(*data_batch)
        image_ids = None

    images = torch.stack(images, 0)
    lengths = torch.tensor([len(cap) for cap in captions], dtype=torch.long)
    captions = pad_sequence(captions, batch_first=True, padding_value=pad_token_id)

    if image_ids is None:
        return images, captions, lengths
    image_ids_tensor = torch.tensor(image_ids, dtype=torch.long)
    return images, captions, lengths, image_ids_tensor


def get_loader(
    transform,
    mode="train",
    batch_size=1,
    tokenizer_name="bert-base-uncased",
    max_length=50,
    num_workers=0,
    cocoapi_loc="/opt",
    subset_size=None,
    subset_fraction=None,
    length_sampler=True,
    shuffle=True,
    return_image_id=False,
    distributed=False,
    dist_rank=0,
    dist_world_size=1,
):
    """Returns the data loader with HuggingFace tokenizer support.
    
    Args:
      transform: Image transform.
      mode: One of 'train', 'val', or 'test'.
      batch_size: Batch size (if in testing mode, must have batch_size=1).
      tokenizer_name: Name of the HuggingFace tokenizer to use (e.g., 'bert-base-uncased').
      max_length: Maximum length for tokenized captions.
      num_workers: Number of subprocesses to use for data loading.
      cocoapi_loc: The location of the folder containing the COCO API: https://github.com/cocodataset/cocoapi
      subset_size: Optional hard limit on the number of samples to iterate over.
      subset_fraction: Fraction (0, 1] of samples to draw from the dataset.
      length_sampler: Whether to reproduce the legacy length-based batching behavior.
      shuffle: Whether to shuffle data when length sampling is disabled.
      return_image_id: Whether to include the COCO image ID in each sample (useful for metrics).
    """

    assert mode in ["train", "val", "test"], "mode must be one of 'train', 'val', or 'test'."
    if subset_size is not None and subset_size <= 0:
        raise ValueError("subset_size must be a positive integer when provided.")
    if subset_fraction is not None:
        if not (0 < subset_fraction <= 1):
            raise ValueError("subset_fraction must be within (0, 1].")
        if subset_size is not None:
            raise ValueError("subset_fraction and subset_size cannot be combined.")

    # Based on mode (train, val, test), obtain img_folder and annotations_file.
    if mode == "train":
        img_folder = os.path.join(cocoapi_loc, "images/train2014/")
        annotations_file = os.path.join(
            cocoapi_loc, "annotations/captions_train2014.json"
        )
    elif mode == "val":
        img_folder = os.path.join(cocoapi_loc, "images/val2014/")
        annotations_file = os.path.join(
            cocoapi_loc, "annotations/captions_val2014.json"
        )
    elif mode == "test":
        assert batch_size == 1, "Please change batch_size to 1 if testing the model."
        img_folder = os.path.join(cocoapi_loc, "images/test2014/")
        annotations_file = os.path.join(
            cocoapi_loc, "annotations/image_info_test2014.json"
        )
    else:
        raise ValueError(f"Invalid mode: {mode}")

    # COCO caption dataset with HuggingFace tokenizer
    dataset = COCODataset(
        transform=transform,
        mode=mode,
        batch_size=batch_size,
        tokenizer_name=tokenizer_name,
        annotations_file=annotations_file,
        img_folder=img_folder,
        max_length=max_length,
        return_image_id=return_image_id,
    )
    
    # Get pad_token_id from the dataset's tokenizer
    pad_token_id = dataset.tokenizer.pad_token_id if dataset.tokenizer.pad_token_id is not None else 0
    
    # Create a collate function with the correct pad_token_id
    def collate_with_padding(batch):
        return collate_fn(batch, pad_token_id=pad_token_id)

    pin_memory = torch.cuda.is_available()
    persistent_workers = num_workers > 0

    if distributed:
        if mode != "train":
            raise ValueError("Distributed sampling is only supported in training mode.")
        if length_sampler:
            raise ValueError("length_sampler cannot be used with distributed sampling.")
        if subset_size is not None or subset_fraction is not None:
            raise ValueError("subset_size/subset_fraction are not supported with distributed sampling.")
        sampler = DistributedSampler(
            dataset,
            num_replicas=dist_world_size,
            rank=dist_rank,
            shuffle=shuffle,
            drop_last=False,
        )
        return data.DataLoader(
            dataset=dataset,
            batch_size=batch_size,
            num_workers=num_workers,
            sampler=sampler,
            collate_fn=collate_with_padding,
            drop_last=False,
            pin_memory=pin_memory,
            persistent_workers=persistent_workers,
        )

    if subset_fraction is not None:
        subset_size = max(1, int(math.ceil(len(dataset) * subset_fraction)))

    if mode == "train" and length_sampler:
        if subset_size is not None:
            raise ValueError("subset_size cannot be combined with length_sampler=True.")
        # Randomly sample a caption length, and sample indices with that length.
        indices = dataset.get_train_indices()
        # Create and assign a batch sampler to retrieve a batch with the sampled indices.
        initial_sampler = data.sampler.SubsetRandomSampler(indices=indices)
        # data loader for COCO dataset.
        data_loader = data.DataLoader(
            dataset=dataset,
            num_workers=num_workers,
            batch_sampler=data.sampler.BatchSampler(
                sampler=initial_sampler, batch_size=dataset.batch_size, drop_last=False
            ),
            collate_fn=collate_with_padding,
            pin_memory=pin_memory,
            persistent_workers=persistent_workers,
        )
    else:
        loader_kwargs = dict(
            dataset=dataset,
            batch_size=dataset.batch_size,
            num_workers=num_workers,
            collate_fn=collate_with_padding,
            drop_last=False,
            pin_memory=pin_memory,
            persistent_workers=persistent_workers,
        )
        sampler = None
        if subset_size is not None:
            subset_size = min(len(dataset), subset_size)
            subset_indices = torch.randperm(len(dataset))[:subset_size].tolist()
            sampler = data.sampler.SubsetRandomSampler(indices=subset_indices)
        if sampler is not None:
            loader_kwargs["sampler"] = sampler
        else:
            # Only expose shuffle flag when we are not overriding the sampler.
            loader_kwargs["shuffle"] = bool(shuffle) if mode == "train" else True

        # Use the requested batch_size (may differ from dataset.batch_size when testing).
        if mode == "train":
            loader_kwargs["batch_size"] = batch_size

        data_loader = data.DataLoader(**loader_kwargs)

    return data_loader