| from __future__ import annotations |
|
|
| import json |
| import random |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| import trackio |
| from model import ContrastiveEncoder, nt_xent, parameter_count |
| from safetensors.torch import save_file |
| from sklearn.linear_model import LogisticRegression |
| from sklearn.metrics import accuracy_score |
| from torch.nn import functional as F |
| from torch.utils.data import DataLoader, TensorDataset |
|
|
| PROJECT_DIR = Path(__file__).resolve().parent |
| ROOT_DIR = PROJECT_DIR.parents[1] |
| DATA_DIR = ROOT_DIR / "projects" / "tiny-vision-foundry" / "data" |
| ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "contrastive-pocket" |
|
|
|
|
| def seed_everything(seed: int) -> None: |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
|
|
|
|
| def load_split(name: str) -> tuple[torch.Tensor, np.ndarray]: |
| frame = pd.read_parquet(DATA_DIR / f"{name}.parquet") |
| pixels = np.stack(frame["image"].to_numpy()).astype(np.float32) / 16.0 |
| labels = frame["label"].to_numpy(dtype=np.int64, copy=True) |
| return torch.from_numpy(pixels.reshape(-1, 1, 8, 8)), labels |
|
|
|
|
| def shift_image(image: torch.Tensor, vertical: int, horizontal: int) -> torch.Tensor: |
| shifted = torch.roll(image, shifts=(vertical, horizontal), dims=(-2, -1)) |
| if vertical > 0: |
| shifted[..., :vertical, :] = 0 |
| elif vertical < 0: |
| shifted[..., vertical:, :] = 0 |
| if horizontal > 0: |
| shifted[..., :, :horizontal] = 0 |
| elif horizontal < 0: |
| shifted[..., :, horizontal:] = 0 |
| return shifted |
|
|
|
|
| def augment_batch(pixels: torch.Tensor) -> torch.Tensor: |
| augmented = pixels.clone() |
| for index in range(len(augmented)): |
| vertical = int(torch.randint(-1, 2, (1,)).item()) |
| horizontal = int(torch.randint(-1, 2, (1,)).item()) |
| augmented[index] = shift_image( |
| augmented[index], |
| vertical, |
| horizontal, |
| ) |
| augmented += torch.randn_like(augmented) * 0.045 |
| dropout_mask = torch.rand_like(augmented) < 0.035 |
| augmented[dropout_mask] = 0 |
| contrast = 0.85 + torch.rand(len(augmented), 1, 1, 1) * 0.3 |
| return torch.clamp(augmented * contrast, 0, 1) |
|
|
|
|
| @torch.inference_mode() |
| def embed(model: ContrastiveEncoder, pixels: torch.Tensor) -> np.ndarray: |
| model.eval() |
| embeddings = [] |
| for start in range(0, len(pixels), 256): |
| embeddings.append(model.encode(pixels[start : start + 256]).numpy()) |
| return np.concatenate(embeddings) |
|
|
|
|
| def label_subset(labels: np.ndarray, per_class: int) -> np.ndarray: |
| selected = [] |
| for label in range(10): |
| selected.extend(np.flatnonzero(labels == label)[:per_class].tolist()) |
| return np.asarray(selected, dtype=np.int64) |
|
|
|
|
| def probe( |
| train_features: np.ndarray, |
| train_labels: np.ndarray, |
| test_features: np.ndarray, |
| test_labels: np.ndarray, |
| per_class: int, |
| ) -> float: |
| indices = label_subset(train_labels, per_class) |
| classifier = LogisticRegression( |
| C=3.0, |
| max_iter=3000, |
| random_state=2035, |
| ) |
| classifier.fit(train_features[indices], train_labels[indices]) |
| return float(accuracy_score(test_labels, classifier.predict(test_features))) |
|
|
|
|
| def probe_suite( |
| train_features: np.ndarray, |
| train_labels: np.ndarray, |
| test_features: np.ndarray, |
| test_labels: np.ndarray, |
| ) -> dict[str, float]: |
| return { |
| str(per_class): probe( |
| train_features, |
| train_labels, |
| test_features, |
| test_labels, |
| per_class, |
| ) |
| for per_class in [10, 25, 100] |
| } |
|
|
|
|
| def invariance_score(model: ContrastiveEncoder, pixels: torch.Tensor) -> dict[str, float]: |
| model.eval() |
| with torch.no_grad(): |
| first = F.normalize(model.encode(augment_batch(pixels[:256])), dim=1) |
| second = F.normalize(model.encode(augment_batch(pixels[:256])), dim=1) |
| positive = float((first * second).sum(dim=1).mean()) |
| mismatched = float((first * second.roll(1, dims=0)).sum(dim=1).mean()) |
| return { |
| "positive_pair_cosine": positive, |
| "mismatched_pair_cosine": mismatched, |
| "invariance_gap": positive - mismatched, |
| } |
|
|
|
|
| def main() -> None: |
| seed_everything(2035) |
| train_pixels, train_labels = load_split("train") |
| test_pixels, test_labels = load_split("test") |
| loader = DataLoader( |
| TensorDataset(train_pixels), |
| batch_size=128, |
| shuffle=True, |
| drop_last=True, |
| generator=torch.Generator().manual_seed(2035), |
| ) |
| random_encoder = ContrastiveEncoder() |
| random_train = embed(random_encoder, train_pixels) |
| random_test = embed(random_encoder, test_pixels) |
| raw_train = train_pixels.reshape(len(train_pixels), -1).numpy() |
| raw_test = test_pixels.reshape(len(test_pixels), -1).numpy() |
|
|
| model = ContrastiveEncoder() |
| optimizer = torch.optim.AdamW(model.parameters(), lr=0.0015, weight_decay=0.002) |
| epochs = 220 |
| scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs) |
| trackio.init( |
| project="contrastive-pocket", |
| name="simclr-8x8-v1", |
| config={ |
| "parameters": parameter_count(model), |
| "epochs": epochs, |
| "temperature": 0.18, |
| "labels_used_in_pretraining": False, |
| }, |
| ) |
| for epoch in range(1, epochs + 1): |
| model.train() |
| losses = [] |
| for (pixels,) in loader: |
| first = model(augment_batch(pixels)) |
| second = model(augment_batch(pixels)) |
| loss = nt_xent(first, second, temperature=0.18) |
| optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| optimizer.step() |
| losses.append(loss.item()) |
| scheduler.step() |
| if epoch == 1 or epoch % 10 == 0: |
| trackio.log( |
| { |
| "epoch": epoch, |
| "contrastive_loss": float(np.mean(losses)), |
| "learning_rate": scheduler.get_last_lr()[0], |
| } |
| ) |
| trackio.finish() |
|
|
| learned_train = embed(model, train_pixels) |
| learned_test = embed(model, test_pixels) |
| results = { |
| "model": "Contrastive Pocket", |
| "parameters": parameter_count(model), |
| "unlabeled_pretraining_examples": len(train_pixels), |
| "epochs": epochs, |
| "embedding_dimensions": learned_train.shape[1], |
| "linear_probe_accuracy_by_examples_per_class": { |
| "contrastive_encoder": probe_suite( |
| learned_train, |
| train_labels, |
| learned_test, |
| test_labels, |
| ), |
| "random_encoder": probe_suite( |
| random_train, |
| train_labels, |
| random_test, |
| test_labels, |
| ), |
| "raw_pixels": probe_suite( |
| raw_train, |
| train_labels, |
| raw_test, |
| test_labels, |
| ), |
| }, |
| "augmentation_invariance_cosine": { |
| "contrastive_encoder": invariance_score(model, train_pixels), |
| "random_encoder": invariance_score(random_encoder, train_pixels), |
| }, |
| "embedding_standard_deviation": float(learned_train.std(axis=0).mean()), |
| } |
| ARTIFACT_DIR.mkdir(parents=True, exist_ok=True) |
| save_file(model.state_dict(), ARTIFACT_DIR / "model.safetensors") |
| (ARTIFACT_DIR / "evaluation.json").write_text( |
| json.dumps(results, indent=2), |
| encoding="utf-8", |
| ) |
| print(json.dumps(results, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
