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