source/train.py

from __future__ import annotations

import json
import random
import shutil
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import trackio
from model import LIFSpikingClassifier, MatchedDenseClassifier, parameter_count
from safetensors.torch import 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_DATA = ROOT_DIR / "projects" / "tiny-vision-foundry" / "data"
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "spike-pocket"
DATA_DIR = PROJECT_DIR / "data"
SEED = 2111


def load_split(name: str, shuffle: bool) -> DataLoader:
    frame = pd.read_parquet(VISION_DATA / f"{name}.parquet")
    pixels = np.stack(frame["image"].to_numpy()).astype(np.float32) / 16
    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=shuffle,
        generator=torch.Generator().manual_seed(SEED),
    )


@torch.inference_mode()
def evaluate(
    model: torch.nn.Module,
    loader: DataLoader,
    *,
    spiking: bool,
    noise: float = 0.0,
) -> dict:
    model.eval()
    correct = 0
    total = 0
    loss = 0.0
    spike_sum = 0.0
    batches = 0
    noise_generator = torch.Generator().manual_seed(SEED + 2)
    spike_generator = torch.Generator().manual_seed(SEED + 3)
    for pixels, labels in loader:
        if noise:
            perturbation = torch.randn(pixels.shape, generator=noise_generator) * noise
            pixels = (pixels + perturbation).clamp(0, 1)
        if spiking:
            logits, spike_rate = model(
                pixels,
                timesteps=32,
                generator=spike_generator,
            )
            spike_sum += float(spike_rate)
            batches += 1
        else:
            logits = model(pixels)
        loss += float(F.cross_entropy(logits, labels, reduction="sum"))
        correct += int((logits.argmax(1) == labels).sum())
        total += len(labels)
    result = {"accuracy": correct / total, "loss": loss / total, "examples": total}
    if spiking:
        result["mean_hidden_spike_rate"] = spike_sum / batches
        result["synaptic_activity_proxy"] = (
            result["mean_hidden_spike_rate"] * model.hidden_dimensions * 32
        )
    return result


def train_model(
    model: torch.nn.Module,
    train_loader: DataLoader,
    validation_loader: DataLoader,
    *,
    spiking: bool,
) -> tuple[dict[str, torch.Tensor], int]:
    optimizer = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=1e-4)
    best = -1.0
    best_epoch = 0
    best_state = None
    for epoch in range(1, 121):
        model.train()
        generator = torch.Generator().manual_seed(SEED + epoch)
        for pixels, labels in train_loader:
            if spiking:
                logits, spike_rate = model(
                    pixels,
                    timesteps=24,
                    generator=generator,
                )
                loss = F.cross_entropy(logits, labels) + 0.002 * spike_rate
            else:
                loss = F.cross_entropy(model(pixels), labels)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
            optimizer.step()
        validation = evaluate(
            model,
            validation_loader,
            spiking=spiking,
        )
        if validation["accuracy"] > best:
            best = validation["accuracy"]
            best_epoch = epoch
            best_state = {
                name: value.detach().cpu().clone()
                for name, value in model.state_dict().items()
            }
        if epoch == 1 or epoch % 10 == 0:
            trackio.log(
                {
                    "variant": "spiking" if spiking else "dense",
                    "epoch": epoch,
                    "validation_accuracy": validation["accuracy"],
                    "validation_loss": validation["loss"],
                    "spike_rate": validation.get("mean_hidden_spike_rate", 0.0),
                }
            )
    assert best_state is not None
    return best_state, best_epoch


def main() -> None:
    random.seed(SEED)
    np.random.seed(SEED)
    torch.manual_seed(SEED)
    torch.set_num_threads(1)
    train_loader = load_split("train", True)
    validation_loader = load_split("validation", False)
    test_loader = load_split("test", False)
    spiking = LIFSpikingClassifier()
    dense = MatchedDenseClassifier()
    trackio.init(
        project="spike-pocket",
        name="surrogate-gradient-lif-v1",
        config={
            "parameters": parameter_count(spiking),
            "training_timesteps": 24,
            "evaluation_timesteps": 32,
            "encoding": "Poisson rate coding",
        },
    )
    spiking_state, spiking_epoch = train_model(
        spiking, train_loader, validation_loader, spiking=True
    )
    dense_state, dense_epoch = train_model(
        dense, train_loader, validation_loader, spiking=False
    )
    spiking.load_state_dict(spiking_state)
    dense.load_state_dict(dense_state)
    results = {
        "spiking_lif": {
            "parameters": parameter_count(spiking),
            "best_epoch": spiking_epoch,
            "clean": evaluate(spiking, test_loader, spiking=True),
            "gaussian_noise_0.20": evaluate(
                spiking, test_loader, spiking=True, noise=0.20
            ),
        },
        "matched_dense": {
            "parameters": parameter_count(dense),
            "best_epoch": dense_epoch,
            "clean": evaluate(dense, test_loader, spiking=False),
            "gaussian_noise_0.20": evaluate(
                dense, test_loader, spiking=False, noise=0.20
            ),
        },
    }
    report = {
        "experiment": "Surrogate-gradient LIF spiking classifier",
        "results": results,
        "activity_boundary": (
            "Spike rate is an activity proxy, not measured hardware energy."
        ),
    }
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    save_file(spiking.state_dict(), ARTIFACT_DIR / "spiking_lif.safetensors")
    save_file(dense.state_dict(), ARTIFACT_DIR / "matched_dense.safetensors")
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(report, indent=2), encoding="utf-8"
    )
    shutil.copy2(VISION_DATA / "test.parquet", DATA_DIR / "test.parquet")
    trackio.log(
        {
            "spiking_clean_accuracy": results["spiking_lif"]["clean"]["accuracy"],
            "dense_clean_accuracy": results["matched_dense"]["clean"]["accuracy"],
            "spike_rate": results["spiking_lif"]["clean"][
                "mean_hidden_spike_rate"
            ],
        }
    )
    trackio.finish()
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()