From 7fd707f6068da18ea1b88100892613937b060936 Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Wed, 22 Jul 2026 05:57:30 -0500 Subject: oral-a: add deterministic CIFAR image pipeline --- sdil/data.py | 123 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 123 insertions(+) (limited to 'sdil') diff --git a/sdil/data.py b/sdil/data.py index abaabe5..7f49b4e 100644 --- a/sdil/data.py +++ b/sdil/data.py @@ -65,6 +65,63 @@ class _FastLoader: return (self.n + self.bs - 1) // self.bs +class _CIFARImageLoader(_FastLoader): + """Deterministic in-memory CIFAR iterator with standard train augmentation. + + Images are stored normalized in NCHW order. Because torchvision applies + zero padding *before* normalization, the normalized padding value is + channel dependent; filling the padded tensor explicitly preserves that + convention instead of padding normalized images with their channel mean. + All shuffle/crop/flip draws use the loader's private CPU generator, so two + loaders with the same seed produce bitwise identical epoch streams. + """ + + def __init__(self, x, y, batch_size, shuffle, augment=False, seed=0): + super().__init__(x, y, batch_size, shuffle, seed=seed) + if x.ndim != 4 or tuple(x.shape[1:]) != (3, 32, 32): + raise ValueError( + f"CIFAR image loader expects NCHW (N,3,32,32), got {tuple(x.shape)}") + if augment and not shuffle: + raise ValueError("CIFAR augmentation is restricted to the shuffled training loader") + self.augment = bool(augment) + self.pad = 4 + + def _augment(self, x): + batch = x.shape[0] + if batch == 0: + return x + # Equivalent to RandomCrop(32, padding=4, fill=0) before Normalize. + fill = torch.as_tensor( + -_CIFAR_MEAN / _CIFAR_STD, dtype=x.dtype, device=x.device) + padded = x.new_empty((batch, 3, 40, 40)) + padded[:] = fill.view(1, 3, 1, 1) + padded[:, :, self.pad:self.pad + 32, self.pad:self.pad + 32] = x + + top = torch.randint(0, 2 * self.pad + 1, (batch,), generator=self.g) + left = torch.randint(0, 2 * self.pad + 1, (batch,), generator=self.g) + flip = torch.rand(batch, generator=self.g) < 0.5 + top = top.to(x.device) + left = left.to(x.device) + flip = flip.to(x.device) + + # Advanced indexing extracts a distinct 32x32 crop for every example. + rows = top[:, None] + torch.arange(32, device=x.device)[None, :] + cols = left[:, None] + torch.arange(32, device=x.device)[None, :] + bhwc = padded.permute(0, 2, 3, 1) + bidx = torch.arange(batch, device=x.device)[:, None, None] + cropped = bhwc[bidx, rows[:, :, None], cols[:, None, :], :] + cropped = cropped.permute(0, 3, 1, 2).contiguous() + return torch.where(flip[:, None, None, None], cropped.flip(-1), cropped) + + def __iter__(self): + idx = (torch.randperm(self.n, generator=self.g) + if self.shuffle else torch.arange(self.n)) + for i in range(0, self.n, self.bs): + j = idx[i:i + self.bs] + x, y = self.x[j], self.y[j] + yield self._augment(x) if self.augment else x, y + + def _load_cifar(data_dir, n_classes=10): root = f"{data_dir}/cifar-10-batches-py" def unpickle(fn): @@ -86,6 +143,72 @@ def _load_cifar(data_dir, n_classes=10): torch.from_numpy(norm(xte)), torch.from_numpy(yte)) +def get_cifar_image_splits(batch_size=128, data_dir=DATA_DIR, device="cpu", + train_limit=None, val_examples=5000, split_seed=2027, + loader_seed=0, augment_train=True): + """Load CIFAR-10 as NCHW images for convolutional protocols. + + Returns ``(train, validation, test, input_shape, n_out, metadata)``. The + validation split is stratified and constructed only from the original + training set. Training augmentation is the canonical random 32x32 crop + after four-pixel zero padding followed by a 50% horizontal flip; validation + and test images are never augmented. This is intentionally a separate API + so all frozen flattened-CIFAR experiments retain their historical tensors. + """ + xtr, ytr, xte, yte = _load_cifar(data_dir) + xtr = xtr.view(-1, 3, 32, 32) + xte = xte.view(-1, 3, 32, 32) + + if val_examples: + train_idx, val_idx = _stratified_validation_indices( + ytr, val_examples, split_seed) + xval, yval = xtr[val_idx], ytr[val_idx] + xtr, ytr = xtr[train_idx], ytr[train_idx] + digest = hashlib.sha256(val_idx.numpy().tobytes()).hexdigest() + val_counts = {str(int(cls)): int((yval == cls).sum()) + for cls in torch.unique(yval, sorted=True)} + else: + xval = yval = None + digest = None + val_counts = {} + if train_limit is not None: + if train_limit <= 0 or train_limit > xtr.shape[0]: + raise ValueError( + f"train_limit must be in [1, {xtr.shape[0]}], got {train_limit}") + xtr, ytr = xtr[:train_limit], ytr[:train_limit] + + xtr, ytr = xtr.to(device), ytr.to(device) + xte, yte = xte.to(device), yte.to(device) + validation = None + if xval is not None: + xval, yval = xval.to(device), yval.to(device) + validation = _CIFARImageLoader( + xval, yval, 1000, False, augment=False, seed=loader_seed) + metadata = { + "dataset": "cifar10", + "input_layout": "NCHW", + "input_shape": [3, 32, 32], + "normalization_mean": _CIFAR_MEAN.tolist(), + "normalization_std": _CIFAR_STD.tolist(), + "training_augmentation": ( + "random_crop_32_padding_4_zero_then_horizontal_flip_p0.5" + if augment_train else "none"), + "loader_seed": int(loader_seed), + "split_seed": int(split_seed) if val_examples else None, + "validation_examples": int(val_examples), + "validation_index_sha256": digest, + "validation_class_counts": val_counts, + "train_examples": int(xtr.shape[0]), + "test_examples": int(xte.shape[0]), + "split_from_training_only": True, + } + train = _CIFARImageLoader( + xtr, ytr, batch_size, True, augment=augment_train, seed=loader_seed) + test = _CIFARImageLoader( + xte, yte, 1000, False, augment=False, seed=loader_seed) + return train, validation, test, (3, 32, 32), 10, metadata + + def _stratified_validation_indices(y, n_val, seed): """Return deterministic disjoint train/validation indices. -- cgit v1.2.3