source/train.py

from __future__ import annotations

import copy
import json
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import trackio
from model import JEncoder, JPredictor, parameter_count
from safetensors.torch import save_file
from sklearn.datasets import load_digits
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader, TensorDataset

PROJECT_DIR = Path(__file__).resolve().parent
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "pocket-jepa"
DATA_DIR = PROJECT_DIR / "data"


def seed_everything(seed: int) -> None:
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.set_num_threads(1)


def mask_blocks(images: torch.Tensor, rng: np.random.Generator) -> torch.Tensor:
    masked = images.clone()
    for index in range(len(masked)):
        size = int(rng.integers(2, 5))
        row = int(rng.integers(0, 9 - size))
        column = int(rng.integers(0, 9 - size))
        masked[index, row : row + size, column : column + size] = 0
    masked += 0.035 * torch.randn_like(masked)
    return masked.clamp(0, 1)


def deterministic_masks(images: np.ndarray, seed: int) -> np.ndarray:
    rng = np.random.default_rng(seed)
    output = images.copy()
    for index in range(len(output)):
        size = 3
        row = int(rng.integers(0, 6))
        column = int(rng.integers(0, 6))
        output[index, row : row + size, column : column + size] = 0
    return output


def redundancy_loss(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    prediction = (prediction - prediction.mean(0)) / (
        prediction.std(0, unbiased=False) + 1e-4
    )
    target = (target - target.mean(0)) / (
        target.std(0, unbiased=False) + 1e-4
    )
    correlation = prediction.T @ target / len(prediction)
    diagonal = torch.diagonal(correlation)
    identity = (diagonal - 1).pow(2).mean()
    off_diagonal = correlation - torch.diag(diagonal)
    return identity + 0.01 * off_diagonal.pow(2).sum() / prediction.shape[1]


@torch.inference_mode()
def encode(model: nn.Module, images: np.ndarray) -> np.ndarray:
    model.eval()
    tensor = torch.from_numpy(images.astype(np.float32))
    embeddings = model(tensor).numpy()
    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
    return embeddings / np.maximum(norms, 1e-8)


def label_budget_indices(
    labels: np.ndarray, per_class: int, seed: int
) -> np.ndarray:
    rng = np.random.default_rng(seed)
    selected = []
    for label in np.unique(labels):
        candidates = np.flatnonzero(labels == label)
        selected.extend(rng.choice(candidates, per_class, replace=False))
    return np.asarray(selected)


def probe(
    train_embeddings: np.ndarray,
    train_labels: np.ndarray,
    test_embeddings: np.ndarray,
    test_labels: np.ndarray,
    selected: np.ndarray,
) -> float:
    classifier = LogisticRegression(C=3.0, max_iter=2_000)
    classifier.fit(train_embeddings[selected], train_labels[selected])
    return float(
        accuracy_score(test_labels, classifier.predict(test_embeddings))
    )


def main() -> None:
    seed_everything(2043)
    digits = load_digits()
    images = (digits.images / 16.0).astype(np.float32)
    labels = digits.target.astype(np.int64)
    indices = np.arange(len(images))
    train_indices, test_indices = train_test_split(
        indices, test_size=0.25, random_state=2043, stratify=labels
    )
    train_images = images[train_indices]
    test_images = images[test_indices]
    train_labels = labels[train_indices]
    test_labels = labels[test_indices]

    online = JEncoder()
    random_control = copy.deepcopy(online)
    target = copy.deepcopy(online)
    predictor = JPredictor()
    for parameter in target.parameters():
        parameter.requires_grad = False
    optimizer = torch.optim.AdamW(
        [*online.parameters(), *predictor.parameters()],
        lr=2e-3,
        weight_decay=2e-4,
    )
    loader = DataLoader(
        TensorDataset(torch.from_numpy(train_images)),
        batch_size=256,
        shuffle=True,
        drop_last=True,
        generator=torch.Generator().manual_seed(2043),
    )
    epochs = 260
    trackio.init(
        project="pocket-jepa",
        name="masked-j-space-v1",
        config={
            "encoder_parameters": parameter_count(online),
            "predictor_parameters": parameter_count(predictor),
            "epochs": epochs,
            "labels_per_class_for_probe": 10,
            "target_ema": 0.99,
        },
    )
    rng = np.random.default_rng(2043)
    history = []
    online.train()
    predictor.train()
    for epoch in range(1, epochs + 1):
        losses = []
        for (complete,) in loader:
            context = mask_blocks(complete, rng)
            prediction = predictor(online(context))
            with torch.no_grad():
                target_embedding = target(complete)
            cosine = 1 - F.cosine_similarity(
                prediction, target_embedding, dim=1
            ).mean()
            decorrelation = redundancy_loss(prediction, target_embedding)
            loss = cosine + 0.35 * decorrelation
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(
                [*online.parameters(), *predictor.parameters()], 1.0
            )
            optimizer.step()
            with torch.no_grad():
                for target_parameter, online_parameter in zip(
                    target.parameters(), online.parameters(), strict=True
                ):
                    target_parameter.mul_(0.99).add_(
                        online_parameter, alpha=0.01
                    )
            losses.append(float(loss.detach()))
        record = {"epoch": epoch, "pretraining_loss": float(np.mean(losses))}
        history.append(record)
        if epoch % 20 == 0:
            trackio.log(record)

    selected = label_budget_indices(train_labels, per_class=10, seed=3043)
    learned_train = encode(target, train_images)
    learned_test = encode(target, test_images)
    learned_masked = encode(target, deterministic_masks(test_images, 4043))
    random_train = encode(random_control, train_images)
    random_test = encode(random_control, test_images)
    random_masked = encode(
        random_control, deterministic_masks(test_images, 4043)
    )
    raw_train = train_images.reshape(len(train_images), -1)
    raw_test = test_images.reshape(len(test_images), -1)
    raw_masked = deterministic_masks(test_images, 4043).reshape(
        len(test_images), -1
    )
    results = {
        "model": "Pocket JEPA",
        "encoder_parameters": parameter_count(target),
        "predictor_training_parameters": parameter_count(predictor),
        "unlabeled_pretraining_examples": len(train_images),
        "pretraining_epochs": epochs,
        "linear_probe_labels": int(len(selected)),
        "labels_per_class": 10,
        "clean_accuracy": {
            "pocket_jepa": probe(
                learned_train,
                train_labels,
                learned_test,
                test_labels,
                selected,
            ),
            "random_encoder": probe(
                random_train,
                train_labels,
                random_test,
                test_labels,
                selected,
            ),
            "raw_pixels": probe(
                raw_train, train_labels, raw_test, test_labels, selected
            ),
        },
        "masked_accuracy": {
            "pocket_jepa": probe(
                learned_train,
                train_labels,
                learned_masked,
                test_labels,
                selected,
            ),
            "random_encoder": probe(
                random_train,
                train_labels,
                random_masked,
                test_labels,
                selected,
            ),
            "raw_pixels": probe(
                raw_train, train_labels, raw_masked, test_labels, selected
            ),
        },
        "final_pretraining_loss": history[-1]["pretraining_loss"],
    }
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    save_file(target.state_dict(), ARTIFACT_DIR / "model.safetensors")
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(results, indent=2), encoding="utf-8"
    )
    np.savez_compressed(
        ARTIFACT_DIR / "j_space.npz",
        embeddings=learned_test,
        images=test_images,
        labels=test_labels,
    )
    pd.DataFrame(
        {
            "source_index": indices,
            "label": labels,
            "split": np.where(
                np.isin(indices, test_indices), "test", "unlabeled_train"
            ),
        }
    ).to_parquet(DATA_DIR / "split_manifest.parquet", index=False)
    trackio.log(
        {
            "clean_probe_accuracy": results["clean_accuracy"]["pocket_jepa"],
            "masked_probe_accuracy": results["masked_accuracy"]["pocket_jepa"],
            "random_clean_accuracy": results["clean_accuracy"][
                "random_encoder"
            ],
        }
    )
    trackio.finish()
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()