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 (
    ConditionalCodePrior,
    TinyVisionJudge,
    VectorQuantizedAutoencoder,
    parameter_count,
)
from PIL import Image
from safetensors.torch import load_file, save_file
from sklearn.cluster import KMeans
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]
VISION_DIR = ROOT_DIR / "projects" / "tiny-vision-foundry"
JUDGE_WEIGHTS = (
    VISION_DIR / "artifacts" / "tiny-student-scratch" / "model.safetensors"
)
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "vq-pocket"
DATA_DIR = PROJECT_DIR / "data"
SEED = 2053


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


def load_split(name: str, *, shuffle: bool, batch_size: int) -> DataLoader:
    frame = pd.read_parquet(VISION_DIR / "data" / 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 DataLoader(
        TensorDataset(
            torch.from_numpy(pixels).reshape(-1, 1, 8, 8),
            torch.from_numpy(labels),
        ),
        batch_size=batch_size,
        shuffle=shuffle,
        generator=torch.Generator().manual_seed(SEED),
    )


@torch.inference_mode()
def collect_codes(
    model: VectorQuantizedAutoencoder,
    loader: DataLoader,
) -> tuple[torch.Tensor, torch.Tensor]:
    model.eval()
    all_codes = []
    all_labels = []
    for pixels, labels in loader:
        _, codes = model.quantize(model.encode(pixels))
        all_codes.append(codes.flatten(1))
        all_labels.append(labels)
    return torch.cat(all_codes), torch.cat(all_labels)


@torch.inference_mode()
def reconstruction_metrics(
    model: VectorQuantizedAutoencoder,
    judge: TinyVisionJudge,
    loader: DataLoader,
) -> dict:
    model.eval()
    judge.eval()
    squared_error = 0.0
    correct = 0
    examples = 0
    code_counts = torch.zeros(model.codebook_size)
    for pixels, labels in loader:
        reconstruction, codes, _, _ = model(pixels)
        squared_error += F.mse_loss(reconstruction, pixels, reduction="sum").item()
        predictions = judge(reconstruction).argmax(dim=1)
        correct += int((predictions == labels).sum())
        examples += len(labels)
        code_counts += torch.bincount(codes.flatten(), minlength=model.codebook_size)
    probabilities = code_counts / code_counts.sum()
    entropy = -(probabilities[probabilities > 0] * probabilities[probabilities > 0].log())
    return {
        "pixel_mse": squared_error / (examples * 64),
        "judge_accuracy": correct / examples,
        "active_codes": int((code_counts > 0).sum()),
        "codebook_size": model.codebook_size,
        "codebook_perplexity": float(entropy.sum().exp()),
        "examples": examples,
    }


@torch.inference_mode()
def generation_metrics(
    autoencoder: VectorQuantizedAutoencoder,
    prior: ConditionalCodePrior,
    judge: TinyVisionJudge,
) -> tuple[dict, torch.Tensor, torch.Tensor, torch.Tensor]:
    autoencoder.eval()
    prior.eval()
    labels = torch.arange(10).repeat_interleave(100)
    codes = prior.generate(labels, seed=SEED + 10_000, temperature=0.85)
    generated = autoencoder.decode_indices(codes)
    predictions = judge(generated).argmax(dim=1)
    per_class = {
        str(label): float(
            (predictions[labels == label] == labels[labels == label]).float().mean()
        )
        for label in range(10)
    }
    diversity = {
        str(label): float(generated[labels == label].flatten(1).var(dim=0).mean())
        for label in range(10)
    }
    flat_codes = codes.flatten(1).numpy()
    unique_by_class = {
        str(label): float(
            len({row.tobytes() for row in flat_codes[labels.numpy() == label]}) / 100
        )
        for label in range(10)
    }
    report = {
        "judge_accuracy": float((predictions == labels).float().mean()),
        "judge_accuracy_by_class": per_class,
        "mean_pixel_variance_by_class": diversity,
        "unique_code_sequence_fraction_by_class": unique_by_class,
        "mean_unique_code_sequence_fraction": float(
            np.mean(list(unique_by_class.values()))
        ),
        "samples": len(labels),
        "sampling_temperature": 0.85,
    }
    return report, generated, labels, codes


def save_grid(generated: torch.Tensor, labels: torch.Tensor, path: Path) -> None:
    images = torch.cat(
        [generated[labels == label][:10] for label in range(10)]
    ).reshape(10, 10, 8, 8)
    canvas = np.zeros((80, 80), dtype=np.uint8)
    for row in range(10):
        for column in range(10):
            canvas[row * 8 : (row + 1) * 8, column * 8 : (column + 1) * 8] = (
                images[row, column].mul(255).clamp(0, 255).to(torch.uint8).numpy()
            )
    Image.fromarray(canvas, mode="L").resize((800, 800), Image.Resampling.NEAREST).save(
        path
    )


def main() -> None:
    seed_everything(SEED)
    torch.set_num_threads(1)
    train_loader = load_split("train", shuffle=True, batch_size=128)
    train_ordered = load_split("train", shuffle=False, batch_size=256)
    validation_loader = load_split("validation", shuffle=False, batch_size=256)
    test_loader = load_split("test", shuffle=False, batch_size=256)
    autoencoder = VectorQuantizedAutoencoder()
    judge = TinyVisionJudge()
    judge.load_state_dict(load_file(JUDGE_WEIGHTS))
    judge.eval()
    for parameter in judge.parameters():
        parameter.requires_grad_(False)
    optimizer = torch.optim.AdamW(autoencoder.parameters(), lr=2e-3, weight_decay=1e-5)
    trackio.init(
        project="vq-pocket",
        name="vq-vae-discrete-prior-v1",
        config={
            "autoencoder_parameters": parameter_count(autoencoder),
            "codebook_size": autoencoder.codebook_size,
            "latent_tokens": 16,
            "continuous_warmup_epochs": 30,
            "vq_epochs": 100,
        },
    )

    for epoch in range(1, 31):
        autoencoder.train()
        running = 0.0
        for pixels, _ in train_loader:
            reconstruction = autoencoder.decode(autoencoder.encode(pixels))
            loss = F.mse_loss(reconstruction, pixels)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
            running += float(loss.detach()) * len(pixels)
        if epoch == 1 or epoch % 10 == 0:
            trackio.log({"phase": 0, "epoch": epoch, "continuous_mse": running / 1257})

    with torch.inference_mode():
        encoded = torch.cat(
            [
                autoencoder.encode(pixels).permute(0, 2, 3, 1).reshape(-1, 16)
                for pixels, _ in train_ordered
            ]
        ).numpy()
    clusters = KMeans(
        n_clusters=autoencoder.codebook_size,
        random_state=SEED,
        n_init=10,
    ).fit(encoded)
    autoencoder.codebook.weight.data.copy_(
        torch.from_numpy(clusters.cluster_centers_).float()
    )
    optimizer = torch.optim.AdamW(autoencoder.parameters(), lr=1e-3, weight_decay=1e-5)
    best_validation = float("inf")
    best_epoch = 0
    best_state = None
    for epoch in range(1, 101):
        autoencoder.train()
        running = 0.0
        for pixels, _ in train_loader:
            reconstruction, _, codebook_loss, commitment_loss = autoencoder(pixels)
            reconstruction_loss = F.mse_loss(reconstruction, pixels)
            loss = reconstruction_loss + codebook_loss + 0.25 * commitment_loss
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
            running += float(loss.detach()) * len(pixels)
        metrics = reconstruction_metrics(autoencoder, judge, validation_loader)
        if metrics["pixel_mse"] < best_validation:
            best_validation = metrics["pixel_mse"]
            best_epoch = epoch
            best_state = {
                name: value.detach().cpu().clone()
                for name, value in autoencoder.state_dict().items()
            }
        if epoch == 1 or epoch % 10 == 0:
            trackio.log(
                {
                    "phase": 1,
                    "epoch": epoch,
                    "vq_loss": running / 1257,
                    "validation_mse": metrics["pixel_mse"],
                    "validation_judge_accuracy": metrics["judge_accuracy"],
                    "active_codes": metrics["active_codes"],
                    "codebook_perplexity": metrics["codebook_perplexity"],
                }
            )
    assert best_state is not None
    autoencoder.load_state_dict(best_state)

    codes, code_labels = collect_codes(autoencoder, train_ordered)
    prior = ConditionalCodePrior(codebook_size=autoencoder.codebook_size)
    prior_optimizer = torch.optim.AdamW(prior.parameters(), lr=2e-3, weight_decay=1e-4)
    prior_dataset = TensorDataset(codes, code_labels)
    prior_loader = DataLoader(
        prior_dataset,
        batch_size=128,
        shuffle=True,
        generator=torch.Generator().manual_seed(SEED),
    )
    start = torch.full((len(codes), 1), prior.start_token, dtype=torch.long)
    best_prior_loss = float("inf")
    best_prior_epoch = 0
    best_prior_state = None
    for epoch in range(1, 181):
        prior.train()
        running = 0.0
        examples = 0
        for batch_codes, batch_labels in prior_loader:
            inputs = torch.cat(
                [
                    torch.full(
                        (len(batch_codes), 1),
                        prior.start_token,
                        dtype=torch.long,
                    ),
                    batch_codes[:, :-1],
                ],
                dim=1,
            )
            logits = prior(inputs, batch_labels)
            loss = F.cross_entropy(
                logits.reshape(-1, prior.codebook_size),
                batch_codes.reshape(-1),
            )
            prior_optimizer.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(prior.parameters(), 5.0)
            prior_optimizer.step()
            running += float(loss.detach()) * len(batch_codes)
            examples += len(batch_codes)
        epoch_loss = running / examples
        if epoch_loss < best_prior_loss:
            best_prior_loss = epoch_loss
            best_prior_epoch = epoch
            best_prior_state = {
                name: value.detach().cpu().clone()
                for name, value in prior.state_dict().items()
            }
        if epoch == 1 or epoch % 20 == 0:
            trackio.log({"phase": 2, "epoch": epoch, "prior_nll": epoch_loss})
    del start
    assert best_prior_state is not None
    prior.load_state_dict(best_prior_state)

    reconstruction = reconstruction_metrics(autoencoder, judge, test_loader)
    generation, generated, labels, generated_codes = generation_metrics(
        autoencoder,
        prior,
        judge,
    )
    results = {
        "model": "VQ-Pocket",
        "method": "VQ-VAE with a class-conditional autoregressive latent-token prior",
        "autoencoder_parameters": parameter_count(autoencoder),
        "prior_parameters": parameter_count(prior),
        "codebook_size": autoencoder.codebook_size,
        "tokens_per_image": 16,
        "best_vq_epoch": best_epoch,
        "best_prior_epoch": best_prior_epoch,
        "prior_training_nll": best_prior_loss,
        "reconstruction": reconstruction,
        "generation": generation,
        "judge": "Frozen Tiny Vision student, 98.52% real-image test accuracy",
    }
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    save_file(autoencoder.state_dict(), ARTIFACT_DIR / "vq_vae.safetensors")
    save_file(prior.state_dict(), ARTIFACT_DIR / "code_prior.safetensors")
    save_grid(generated, labels, ARTIFACT_DIR / "samples.png")
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(results, indent=2),
        encoding="utf-8",
    )
    pd.DataFrame(
        {
            "label": labels.numpy(),
            "codes": list(generated_codes.flatten(1).numpy()),
            "pixels": list(generated.flatten(1).numpy()),
        }
    ).to_parquet(DATA_DIR / "generated_token_sequences.parquet", index=False)
    trackio.log(
        {
            "test_reconstruction_mse": reconstruction["pixel_mse"],
            "test_reconstruction_judge_accuracy": reconstruction["judge_accuracy"],
            "generation_judge_accuracy": generation["judge_accuracy"],
            "generation_unique_sequences": generation[
                "mean_unique_code_sequence_fraction"
            ],
        }
    )
    trackio.finish()
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()