source/train.py

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()