oral-a: add deterministic CIFAR image pipeline

2 files changed, 212 insertions, 0 deletions

diff --git a/experiments/cifar_image_smoke.py b/experiments/cifar_image_smoke.py
new file mode 100644
index 0000000..c23da45
--- /dev/null
+++ b/experiments/cifar_image_smoke.py
@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""Deterministic checks for the convolutional CIFAR data path."""
+import argparse
+import os
+import sys
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from sdil.data import _CIFARImageLoader, get_cifar_image_splits
+
+
+EXPECTED_SPLIT_100_SHA256 = (
+    "a88b2ca885bba36dc7407171e3afbce8dca2dd83c7c1d4222e4418b0ca3de6f0")
+
+
+def synthetic_checks():
+    x = (torch.arange(12 * 3 * 32 * 32, dtype=torch.float32)
+         .reshape(12, 3, 32, 32) / 1000.0)
+    y = torch.arange(12)
+    first = _CIFARImageLoader(x, y, 5, True, augment=True, seed=17)
+    second = _CIFARImageLoader(x.clone(), y.clone(), 5, True, augment=True, seed=17)
+    epoch_sums = []
+    for _ in range(2):
+        stream_a, stream_b = list(first), list(second)
+        assert len(stream_a) == len(stream_b) == 3
+        assert all(torch.equal(xa, xb) and torch.equal(ya, yb)
+                   for (xa, ya), (xb, yb) in zip(stream_a, stream_b))
+        assert sum(len(labels) for _, labels in stream_a) == 12
+        assert all(tuple(images.shape[1:]) == (3, 32, 32)
+                   for images, _ in stream_a)
+        epoch_sums.append(sum(float(images.sum()) for images, _ in stream_a))
+    assert epoch_sums[0] != epoch_sums[1], "augmentation RNG did not advance"
+
+    plain_a = _CIFARImageLoader(x, y, 5, False, augment=False, seed=1)
+    plain_b = _CIFARImageLoader(x, y, 5, False, augment=False, seed=999)
+    assert all(torch.equal(xa, xb) and torch.equal(ya, yb)
+               for (xa, ya), (xb, yb) in zip(plain_a, plain_b))
+    try:
+        _CIFARImageLoader(x, y, 5, False, augment=True)
+    except ValueError:
+        pass
+    else:
+        raise AssertionError("evaluation augmentation guard is missing")
+
+
+def real_data_checks(device):
+    args = dict(
+        batch_size=7, device=device, train_limit=13, val_examples=100,
+        split_seed=2027, loader_seed=11)
+    train, validation, test, shape, n_out, metadata = get_cifar_image_splits(**args)
+    assert shape == (3, 32, 32) and n_out == 10
+    assert tuple(train.x.shape) == (13, 3, 32, 32)
+    assert tuple(validation.x.shape) == (100, 3, 32, 32)
+    assert tuple(test.x.shape) == (10000, 3, 32, 32)
+    assert metadata["validation_class_counts"] == {str(i): 10 for i in range(10)}
+    assert metadata["validation_index_sha256"] == EXPECTED_SPLIT_100_SHA256
+
+    validation_before = validation.x.clone()
+    test_before = test.x[:100].clone()
+    list(validation)
+    list(test)
+    assert torch.equal(validation.x, validation_before)
+    assert torch.equal(test.x[:100], test_before)
+
+    paired, _, _, _, _, paired_metadata = get_cifar_image_splits(**args)
+    for (xa, ya), (xb, yb) in zip(train, paired):
+        assert torch.equal(xa, xb) and torch.equal(ya, yb)
+    assert metadata == paired_metadata
+    return metadata
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--synthetic-only", action="store_true")
+    args = parser.parse_args()
+    synthetic_checks()
+    metadata = None if args.synthetic_only else real_data_checks(args.device)
+    print({
+        "synthetic": "passed",
+        "real_data": "skipped" if metadata is None else "passed",
+        "validation_index_sha256": (
+            None if metadata is None else metadata["validation_index_sha256"]),
+    })
+
+
+if __name__ == "__main__":
+    main()
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.