oral-a: add deterministic CIFAR image pipeline

1 files changed, 123 insertions, 0 deletions

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.