File size: 3,246 Bytes
256c9c2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | """
Denoising Diffusion Probabilistic Models — Data Loading
Paper: https://arxiv.org/abs/2006.11239
Authors: Ho, Jain, Abbeel (2020)
§4 — "We set T = 1000... on CIFAR10 (32×32)... We used random horizontal
flips during training; we tried training both with and without dropout..."
Data preprocessing:
- Normalize images to [-1, 1] range (not [0,1])
[FROM_OFFICIAL_CODE] — Paper does not state normalization range explicitly,
but the official code normalizes to [-1,1] which is consistent with
the Gaussian output distribution N(x_0; μ_θ, σ²I).
- Random horizontal flips (§4)
- No other augmentation mentioned
"""
from typing import Optional, Tuple
import torch
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets, transforms
def get_cifar10_transforms(image_size: int = 32) -> Tuple[transforms.Compose, transforms.Compose]:
"""Return train and test transforms for CIFAR-10.
§4 — "random horizontal flips during training"
[FROM_OFFICIAL_CODE] — normalize to [-1, 1]
Args:
image_size: Target image size (CIFAR-10 is natively 32×32)
Returns:
(train_transform, test_transform)
"""
train_transform = transforms.Compose([
transforms.RandomHorizontalFlip(), # §4 — explicit mention
transforms.ToTensor(), # [0, 255] -> [0, 1]
transforms.Normalize( # [0, 1] -> [-1, 1]
mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5],
),
])
test_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(
mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5],
),
])
return train_transform, test_transform
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
|