Function get_dataloaders

skills/paper2code/worked/ddpm/src/data.py:58–111

Create CIFAR-10 train and test dataloaders. §4 — CIFAR-10 unconditional generation at 32×32 §4 — "batch size 128" [FROM_OFFICIAL_CODE — paper does not state batch size] Args: data_dir: Root directory for dataset download/cache batch_size: Batch size for training

(
    data_dir: str = "./data",
    batch_size: int = 128,
    num_workers: int = 4,
    image_size: int = 32,
)

def get_dataloaders(
    data_dir: str = "./data",
    batch_size: int = 128,
    num_workers: int = 4,
    image_size: int = 32,
) -> Tuple[DataLoader, DataLoader]:
    """Create CIFAR-10 train and test dataloaders.

    §4 — CIFAR-10 unconditional generation at 32×32
    §4 — "batch size 128" [FROM_OFFICIAL_CODE — paper does not state batch size]

    Args:
        data_dir: Root directory for dataset download/cache
        batch_size: Batch size for training
        num_workers: Number of data loading workers
        image_size: Image resolution (32 for CIFAR-10)

    Returns:
        (train_loader, test_loader)
    """
    train_transform, test_transform = get_cifar10_transforms(image_size)

    train_dataset = datasets.CIFAR10(
        root=data_dir,
        train=True,
        download=True,
        transform=train_transform,
    )

    test_dataset = datasets.CIFAR10(
        root=data_dir,
        train=False,
        download=True,
        transform=test_transform,
    )

    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True,
        drop_last=True,  # [ASSUMPTION] drop incomplete last batch for stable training
    )

    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True,
    )

    return train_loader, test_loader