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 PocketDenoiser, TinyVisionJudge, parameter_count
from PIL import Image
from safetensors.torch import load_file, save_file
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"
DATA_DIR = VISION_DIR / "data"
JUDGE_WEIGHTS = VISION_DIR / "artifacts" / "tiny-student-scratch" / "model.safetensors"
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "pocket-diffusion"
STEPS = 50


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


def load_training_data() -> DataLoader:
    frame = pd.read_parquet(DATA_DIR / "train.parquet")
    pixels = np.stack(frame["image"].to_numpy()).astype(np.float32) / 8.0 - 1.0
    labels = frame["label"].to_numpy(dtype=np.int64, copy=True)
    return DataLoader(
        TensorDataset(torch.from_numpy(pixels), torch.from_numpy(labels)),
        batch_size=128,
        shuffle=True,
        generator=torch.Generator().manual_seed(2032),
    )


def schedule() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    betas = torch.linspace(1e-4, 0.025, STEPS)
    alphas = 1.0 - betas
    cumulative = torch.cumprod(alphas, dim=0)
    return betas, alphas, cumulative


@torch.inference_mode()
def sample(
    model: PocketDenoiser,
    labels: torch.Tensor,
    guidance: float,
    seed: int,
) -> torch.Tensor:
    generator = torch.Generator().manual_seed(seed)
    betas, alphas, cumulative = schedule()
    pixels = torch.randn(len(labels), 64, generator=generator)
    null_labels = torch.full_like(labels, 10)
    model.eval()
    for step in reversed(range(STEPS)):
        timesteps = torch.full((len(labels),), step, dtype=torch.long)
        conditional = model(pixels, timesteps, labels)
        unconditional = model(pixels, timesteps, null_labels)
        predicted_noise = unconditional + guidance * (conditional - unconditional)
        alpha = alphas[step]
        cumulative_alpha = cumulative[step]
        mean = (
            pixels - (1 - alpha) / torch.sqrt(1 - cumulative_alpha) * predicted_noise
        ) / torch.sqrt(alpha)
        if step:
            noise = torch.randn(pixels.shape, generator=generator)
            pixels = mean + torch.sqrt(betas[step]) * noise
        else:
            pixels = mean
    return torch.clamp((pixels + 1) / 2, 0, 1)


@torch.inference_mode()
def generation_metrics(
    model: PocketDenoiser,
    judge: TinyVisionJudge,
    guidance: float,
) -> tuple[dict, torch.Tensor, torch.Tensor]:
    labels = torch.arange(10).repeat_interleave(100)
    generated = sample(model, labels, guidance=guidance, seed=2032)
    predictions = judge(generated.reshape(-1, 1, 8, 8)).argmax(dim=1)
    accuracy_by_class = {
        str(label): float(
            (predictions[labels == label] == labels[labels == label]).float().mean()
        )
        for label in range(10)
    }
    diversity = {
        str(label): float(generated[labels == label].var(dim=0).mean())
        for label in range(10)
    }
    return (
        {
            "judge_accuracy": float((predictions == labels).float().mean()),
            "judge_accuracy_by_class": accuracy_by_class,
            "mean_pixel_variance_by_class": diversity,
            "samples": len(labels),
            "guidance": guidance,
        },
        generated,
        labels,
    )


def save_grid(generated: torch.Tensor, labels: torch.Tensor, path: Path) -> None:
    selected = [generated[labels == label][:10] for label in range(10)]
    images = torch.cat(selected).reshape(10, 10, 8, 8).cpu().numpy()
    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,
            ] = np.clip(images[row, column] * 255, 0, 255).astype(np.uint8)
    Image.fromarray(canvas, mode="L").resize((800, 800), Image.Resampling.NEAREST).save(
        path
    )


def main() -> None:
    seed_everything(2032)
    loader = load_training_data()
    model = PocketDenoiser()
    judge = TinyVisionJudge()
    judge.load_state_dict(load_file(JUDGE_WEIGHTS))
    judge.eval()
    betas, _, cumulative = schedule()
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.0015, weight_decay=0.001)
    epochs = 300
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    trackio.init(
        project="pocket-diffusion",
        name="conditional-ddpm-cfg-v1",
        config={
            "parameters": parameter_count(model),
            "diffusion_steps": STEPS,
            "epochs": epochs,
            "label_dropout": 0.12,
        },
    )
    for epoch in range(1, epochs + 1):
        model.train()
        running_loss = 0.0
        examples = 0
        for pixels, labels in loader:
            timesteps = torch.randint(0, STEPS, (len(labels),))
            noise = torch.randn_like(pixels)
            cumulative_alpha = cumulative[timesteps].unsqueeze(1)
            noisy = (
                torch.sqrt(cumulative_alpha) * pixels
                + torch.sqrt(1 - cumulative_alpha) * noise
            )
            conditioned_labels = labels.clone()
            drop = torch.rand(len(labels)) < 0.12
            conditioned_labels[drop] = 10
            prediction = model(noisy, timesteps, conditioned_labels)
            loss = F.mse_loss(prediction, noise)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
            running_loss += loss.item() * len(labels)
            examples += len(labels)
        scheduler.step()
        if epoch == 1 or epoch % 10 == 0:
            trackio.log(
                {
                    "epoch": epoch,
                    "noise_prediction_mse": running_loss / examples,
                    "learning_rate": scheduler.get_last_lr()[0],
                }
            )
    trackio.finish()

    guidance_candidates = {}
    for guidance in [1.0, 1.5, 2.0, 2.5, 3.0]:
        metrics, _, _ = generation_metrics(model, judge, guidance)
        guidance_candidates[str(guidance)] = metrics["judge_accuracy"]
    best_guidance = float(max(guidance_candidates, key=guidance_candidates.get))
    generation, generated, labels = generation_metrics(model, judge, best_guidance)
    results = {
        "model": "PocketDiffusion",
        "parameters": parameter_count(model),
        "diffusion_steps": STEPS,
        "epochs": epochs,
        "guidance_search": guidance_candidates,
        "generation": generation,
        "judge": "Tiny Vision labels-only student, 98.52% real-image test accuracy",
    }
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    save_file(model.state_dict(), ARTIFACT_DIR / "model.safetensors")
    save_grid(generated, labels, ARTIFACT_DIR / "samples.png")
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(results, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()