cifar10dvs_conv_dyned_dynedc_snn_v9.py
CIFAR10-DVS DyNED + DyNEDc SNN - 2,171 lines.
View on GitHub (image-neuro/cifar10dvs_conv_dyned_dynedc_snn_v9.py).
Source
Section titled “Source”"""
CIFAR10-DVS SNN with VGGSNN Conv + DyNED + DyNEDc Compression + TET Loss (v9)
Pipeline (training):
DVS events -> dense frames -> VGG conv -> proj_down -> DyNED encoder
(continuous, STE-quantised) -> proj_up -> cAdLIF SNN
Pipeline (post-training compression measurement):
Best model -> forward to cAdLIF spike outputs (binary spike trains, native
to cAdLIF) -> DyNEDc compress per layer per sample -> stats JSON
DyNEDc-in-the-data-path notes:
- DyNED runs inside the model (end-to-end with the conv front-end) and produces
continuous level-quantised values that feed the cAdLIF stack directly. We do
not binarise between DyNED and cAdLIF: that route was tested in earlier v9
iterations and the in-forward STE through the binarisation step capped
accuracy at ~50% on CIFAR10-DVS.
- The cAdLIF layers themselves emit binary spike trains during inference;
`measure_compression_stats` runs DyNEDc directly on those spike outputs
(`spk1`, `spk2`, `spk3`). This makes the thesis claim "DyNEDc compresses
spike trains" literal: the compressed objects are the actual spike trains
the SNN produces.
- DyNEDc compression itself is lossless. It runs once post-training over the
test set; results are written to dynedc_compression_stats.json.
- Output dir: cifar10dvs_conv_dyned_dynedc_output_v9
Run: uv run python image-neuro/cifar10dvs_conv_dyned_dynedc_snn_v9.py
Optuna: uv run python image-neuro/cifar10dvs_conv_dyned_dynedc_snn_v9.py --optuna --n-trials 50
Quick smoke: uv run python image-neuro/cifar10dvs_conv_dyned_dynedc_snn_v9.py --subset-size 256 --quick-epochs 2
"""
import argparse
import gc
import math
import os
import time
from pathlib import Path
import matplotlib
try:
matplotlib.use("Agg")
except Exception:
pass
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset, random_split
try:
import tonic
import tonic.transforms as tonic_transforms
HAS_TONIC = True
except ImportError:
HAS_TONIC = False
try:
from sklearn.manifold import TSNE
HAS_TSNE = True
except ImportError:
HAS_TSNE = False
import sys as _sys
_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), os.pardir))
if _REPO_ROOT not in _sys.path:
_sys.path.insert(0, _REPO_ROOT)
from dyned import DyNEDcCompressorV4 # noqa: E402
from vis_utils import dump_plot_data # noqa: E402
# Constants and Configurations
try:
_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
except NameError:
_SCRIPT_DIR = os.getcwd()
OUTPUT_DIR = os.path.join(_SCRIPT_DIR, "cifar10dvs_conv_dyned_dynedc_output_v9")
os.makedirs(OUTPUT_DIR, exist_ok=True)
# Module-level CLI overrides (set in __main__ block; default None means "use script defaults")
_CLI_SUBSET_SIZE = None
_CLI_QUICK_EPOCHS = None
_CLI_DYNEDC_IN_FORWARD = False
CIFAR10DVS_CLASSES = [
"airplane",
"automobile",
"bird",
"cat",
"deer",
"dog",
"frog",
"horse",
"ship",
"truck",
]
_SPATIAL_FACTOR = 0.25
_TARGET_H = 32
_TARGET_W = 32
_N_TIME_BINS = 16
_N_CHANNELS = 2 # ON/OFF polarity
# =============================================================================
# CIFAR10-DVS Dataset Wrapper - dual polarity, 16 time bins
# =============================================================================
class CIFAR10DVSFrameDataset(Dataset):
"""Wraps tonic.datasets.CIFAR10DVS and converts events to temporal frames.
Pipeline per sample:
events (t,x,y,p) -> Downsample(0.25): 128x128 -> 32x32
-> ToFrame(n_time_bins=16): [16, 2, 32, 32]
-> normalize per-frame to [-1, 1]
"""
def __init__(self, save_to, augment=False, cache=True):
if not HAS_TONIC:
raise ImportError("tonic is required: uv add tonic")
self.augment = augment
self.cache = cache
self._data = None
self._labels = None
self._class_list = CIFAR10DVS_CLASSES
# Cache lives in this script's own folder, matching the convention used by
# the speech / SSC scripts (assets/<script_basename>/...). The raw
# event-frame preprocessing is identical to v8, so if a v8 cache is
# already on disk we read it from there to avoid a redundant 1.3 GB
# rebuild - but new caches always go to the v9 folder.
cache_path = Path(save_to) / "cifar10dvs_conv_dyned_dynedc_snn_v9" / "cifar10dvs_frames_cache.pt"
v8_fallback = Path(save_to) / "cifar10dvs_conv_dyned_snn_v8" / "cifar10dvs_frames_cache.pt"
if cache:
for source in (cache_path, v8_fallback):
if source.exists():
if source != cache_path:
print(
f"Loading cached CIFAR10-DVS frames from {source} "
"(v8 fallback; raw event-frame preprocessing is identical)..."
)
else:
print(f"Loading cached CIFAR10-DVS frames from {source}...")
cached = torch.load(source, weights_only=True)
self._data = cached["data"]
self._labels = cached["labels"]
print(f"Loaded {len(self._data)} samples ({self._data.shape})")
return
transform = tonic_transforms.Compose(
[
tonic_transforms.Downsample(spatial_factor=_SPATIAL_FACTOR),
tonic_transforms.ToFrame(
sensor_size=(_TARGET_W, _TARGET_H, 2),
n_time_bins=_N_TIME_BINS,
),
]
)
raw_ds = tonic.datasets.CIFAR10DVS(save_to=save_to, transform=transform)
if cache:
self._build_cache(raw_ds, cache_path)
else:
self._raw_ds = raw_ds
def _build_cache(self, raw_ds, cache_path):
print(f"Caching {len(raw_ds)} CIFAR10-DVS samples as temporal frame tensors...")
t0 = time.time()
frames_list = []
labels_list = []
for i in range(len(raw_ds)):
frame_tensor, label = raw_ds[i]
frame_t = torch.from_numpy(frame_tensor).float() # [16, 2, 32, 32]
# Normalize per-frame to [-1, 1] (across both channels)
T_f = frame_t.shape[0]
flat = frame_t.reshape(T_f, -1)
f_min = flat.min(dim=1).values.view(T_f, 1, 1, 1)
f_max = flat.max(dim=1).values.view(T_f, 1, 1, 1)
f_range = (f_max - f_min).clamp(min=1e-8)
frame_t = (frame_t - f_min) / f_range * 2.0 - 1.0
frames_list.append(frame_t)
labels_list.append(label)
if (i + 1) % 1000 == 0:
elapsed = time.time() - t0
rate = (i + 1) / elapsed
remaining = (len(raw_ds) - i - 1) / rate
print(f" {i + 1}/{len(raw_ds)} ({rate:.0f} samples/s, ~{remaining:.0f}s remaining)")
self._data = torch.stack(frames_list) # [N, 16, 2, 32, 32]
self._labels = torch.tensor(labels_list, dtype=torch.long)
elapsed = time.time() - t0
print(f"Cache built in {elapsed:.1f}s - shape={self._data.shape}, dtype={self._data.dtype}")
cache_path.parent.mkdir(parents=True, exist_ok=True)
torch.save({"data": self._data, "labels": self._labels}, cache_path)
print(f"Saved cache to {cache_path}")
def __len__(self):
if self.cache:
return len(self._data)
return len(self._raw_ds)
def __getitem__(self, idx):
if self.cache:
frames = self._data[idx] # [16, 2, 32, 32]
label = self._labels[idx].item()
else:
frame_tensor, label = self._raw_ds[idx]
frame_t = torch.from_numpy(frame_tensor).float()
T_f = frame_t.shape[0]
flat = frame_t.reshape(T_f, -1)
f_min = flat.min(dim=1).values.view(T_f, 1, 1, 1)
f_max = flat.max(dim=1).values.view(T_f, 1, 1, 1)
f_range = (f_max - f_min).clamp(min=1e-8)
frames = (frame_t - f_min) / f_range * 2.0 - 1.0
return frames, label
@staticmethod
def _random_affine(frames):
"""Apply same spatial transform to all temporal frames."""
T = frames.shape[0]
tx = (torch.rand(1).item() - 0.5) * 0.2
ty = (torch.rand(1).item() - 0.5) * 0.2
scale = 0.9 + torch.rand(1).item() * 0.2
theta = (
torch.tensor(
[
[scale, 0.0, tx],
[0.0, scale, ty],
],
dtype=torch.float32,
)
.unsqueeze(0)
.expand(T, -1, -1)
)
grid = F.affine_grid(theta, frames.shape, align_corners=False)
out = F.grid_sample(frames, grid, align_corners=False, padding_mode="zeros")
return out
class _AugmentedSubset(Dataset):
"""Training augmentation: flip, affine, cutout, EventDrop."""
def __init__(self, subset):
self.subset = subset
def __len__(self):
return len(self.subset)
def __getitem__(self, idx):
frames, label = self.subset[idx] # [16, 2, 32, 32]
# Random horizontal flip
if torch.rand(1).item() > 0.5:
frames = frames.flip(-1)
# Random affine
frames = CIFAR10DVSFrameDataset._random_affine(frames)
# Cutout: zero out a random 6x6 patch across all frames (30% prob)
if torch.rand(1).item() < 0.3:
frames = self._cutout(frames, patch_size=6)
# EventDrop: zero out ~10% of time bins (30% prob)
if torch.rand(1).item() < 0.3:
frames = self._event_drop(frames, drop_ratio=0.1)
return frames, label
@staticmethod
def _cutout(frames, patch_size=6):
"""Zero out a random spatial patch across all frames and channels."""
_, _, H, W = frames.shape
cy = torch.randint(0, H, (1,)).item()
cx = torch.randint(0, W, (1,)).item()
y1 = max(0, cy - patch_size // 2)
y2 = min(H, cy + patch_size // 2)
x1 = max(0, cx - patch_size // 2)
x2 = min(W, cx + patch_size // 2)
frames = frames.clone()
frames[:, :, y1:y2, x1:x2] = 0.0
return frames
@staticmethod
def _event_drop(frames, drop_ratio=0.1):
"""Zero out random time bins to simulate event loss."""
T = frames.shape[0]
n_drop = max(1, int(T * drop_ratio))
drop_idx = torch.randperm(T)[:n_drop]
frames = frames.clone()
frames[drop_idx] = 0.0
return frames
# =============================================================================
# In-forward ON/OFF binarisation + DyNEDc round-trip (port of v8's binary path)
# =============================================================================
def ste_round(x):
"""Straight-through-estimator round: forward = round(x), backward = identity."""
return x + (x.round() - x).detach()
def encode_on_off_torch(quantised, levels):
"""Differentiable ON/OFF event encoder for v9's DyNED output.
v9's DyNED encoder runs sigma-delta along the *time* axis (consecutive
timesteps of the same channel have small deltas). The natural axis for
differencing here is therefore time - consecutive channels at a fixed
timestep are unrelated.
Mapping (per (batch, channel) row, along the time axis t):
1. Per-row min/max normalisation: idx[t] = round((q[t] - row_min) /
row_range * (levels - 1)) clipped to [0, levels-1].
2. delta[t] = idx[t] - idx[t-1] for t >= 1.
3. Clamp delta to {-1, 0, +1}.
4. on[t] = (delta == +1) -> ON channel
off[t] = (delta == -1) -> OFF channel
5. First timestep is initial state -> on[0] = off[0] = 0.
Gradients flow via STE on the rounding step; the {-1, 0, +1} clamp and the
>0 / <0 indicator are wrapped so backward sees the underlying continuous
values.
Args:
quantised: tensor of shape [..., T] (continuous DyNED-quantised values).
The last dim is treated as the time axis.
levels: number of DyNED quantisation levels (e.g. 256).
Returns:
on_off: tensor of shape [..., 2, T] with the ON channel at index 0 and
the OFF channel at index 1 along the new -2 axis. Values are in
{0, 1} in forward but carry STE gradients in backward.
"""
# Per-row min/max along the time axis
row_min = quantised.amin(dim=-1, keepdim=True)
row_max = quantised.amax(dim=-1, keepdim=True)
row_range = (row_max - row_min).clamp(min=1e-8)
norm = (quantised - row_min) / row_range
# STE-rounded indices in [0, levels-1]
idx = ste_round(norm * (levels - 1)).clamp(0.0, float(levels - 1))
# Time-axis differencing then clamp to {-1, 0, +1}. The clamp is
# piecewise-linear (already differentiable) and STE-rounding the
# clamped value yields a {0, 1} mask that still carries gradients.
delta = idx[..., 1:] - idx[..., :-1] # [..., T-1]
delta_clamped = delta.clamp(-1.0, 1.0)
on = ste_round(delta_clamped.clamp(min=0.0)) # 1 if delta == +1 else 0
off = ste_round((-delta_clamped).clamp(min=0.0)) # 1 if delta == -1 else 0
# Pad a leading zero timestep so shape matches the original T
pad_shape = list(on.shape)
pad_shape[-1] = 1
pad = torch.zeros(pad_shape, dtype=on.dtype, device=on.device)
on_full = torch.cat([pad, on], dim=-1)
off_full = torch.cat([pad, off], dim=-1)
# Stack along a new channel dim before T
return torch.stack([on_full, off_full], dim=-2) # [..., 2, T]
def dynedc_compress_binary(binary_2d, chunk_size=4):
"""Compress a binary array with DyNEDcCompressorV4.
Mirrors v8's helper of the same name. Returns (bit_string, info_dict,
codec_state) so the per-sample decoder can be primed with the same
Huffman table/mode that the encoder used.
"""
flat = np.asarray(binary_2d).astype(np.uint8).flatten()
compressor = DyNEDcCompressorV4(chunk_size=chunk_size)
compressed, info = compressor.compress(flat)
codec_state = {
"mode": compressor._mode,
"huff_codes": dict(compressor._huff_codes),
"alt_start": compressor._alt_start,
}
return compressed, info, codec_state
def dynedc_decompress_binary(compressed_str, shape, codec_state, chunk_size=4):
"""Inverse of `dynedc_compress_binary`. Restores per-sample compressor
state from `codec_state` so V4's stateful decoder works correctly.
"""
compressor = DyNEDcCompressorV4(chunk_size=chunk_size)
compressor._mode = codec_state["mode"]
compressor._huff_codes = codec_state["huff_codes"]
compressor._alt_start = codec_state["alt_start"]
decompressed_str = compressor.decompress(compressed_str)
flat = np.frombuffer(decompressed_str.encode("ascii"), dtype=np.uint8) - ord("0")
n_expected = int(np.prod(shape))
flat = flat[:n_expected]
return flat.reshape(shape).astype(np.uint8)
def dynedc_round_trip(binary_tensor, chunk_size=4):
"""Compress + decompress per-sample binary tensor through DyNEDcCompressorV4.
Because the codec is lossless, the returned tensor is bit-identical to the
input. The only purpose is to put DyNEDc *into the data path* in a way
that's verifiable (and that lets us measure the per-batch overhead).
The round-trip is non-differentiable (binary -> binary identity); the
returned tensor is grafted onto the input tensor's autograd graph via
`input + (round_trip - input).detach()` by the caller, so gradients still
flow through the upstream binarisation.
Args:
binary_tensor: tensor of shape [B, ...] with values in {0, 1}. The
leading dim is treated as the batch dim and each sample
is compressed independently.
chunk_size: DyNEDc V4 chunk size (default 4).
Returns:
Tensor with the same shape, dtype and device as the input. Values are
identical to the input (lossless).
"""
if binary_tensor.numel() == 0:
return binary_tensor.clone()
orig_dtype = binary_tensor.dtype
orig_device = binary_tensor.device
arr = binary_tensor.detach().to(torch.uint8).cpu().numpy()
out = np.empty_like(arr)
sample_shape = arr.shape[1:]
for i in range(arr.shape[0]):
compressed, _, codec_state = dynedc_compress_binary(arr[i], chunk_size=chunk_size)
out[i] = dynedc_decompress_binary(compressed, sample_shape, codec_state, chunk_size=chunk_size)
return torch.from_numpy(out).to(device=orig_device, dtype=orig_dtype)
# =============================================================================
# DyNEDc Compression Layer (lossless: passthrough during training, stats during eval)
# =============================================================================
def quantise_to_indices(quantised_2d, levels):
"""Convert continuous DyNED-quantised values to per-row uint8 indices.
Used by the post-training compression measurement pass to convert each
sample's DyNED output into the integer-index representation that
DyNEDcCompressorV4 expects.
"""
arr = quantised_2d.detach().cpu().numpy() if hasattr(quantised_2d, "detach") else np.asarray(quantised_2d)
if arr.ndim == 1:
arr = arr.reshape(1, -1)
n_rows, n_cols = arr.shape
row_min = arr.min(axis=1).astype(np.float32)
row_max = arr.max(axis=1).astype(np.float32)
row_range = row_max - row_min
row_range[row_range < 1e-8] = 1.0
norm = (arr - row_min[:, None]) / row_range[:, None]
indices = np.round(norm * (levels - 1)).clip(0, levels - 1).astype(np.uint8)
return indices
def measure_compression_stats(net, test_loader, device, dyned_levels, chunk_size=4):
"""Post-training DyNEDc compression stats on the cAdLIF spike outputs.
The cAdLIF layers natively emit binary spike trains during inference;
DyNEDc compresses those spike trains losslessly, which is exactly the
thesis claim ("DyNEDc compresses spike trains"). Stats are aggregated
over each layer's `[T, B, hidden]` spike tensor, per sample.
"""
from collections import Counter
net.eval()
layer_keys = ("spk1", "spk2", "spk3")
per_layer = {k: {"ratios": [], "modes": []} for k in layer_keys}
with torch.no_grad():
for data, _ in test_loader:
data = data.to(device)
layer_data = net.diagnostic_forward(data)
for k in layer_keys:
spk = layer_data[k] # [T, B, hidden]
spk_np = spk.permute(1, 0, 2).cpu().numpy().astype(np.uint8) # [B, T, hidden]
B = spk_np.shape[0]
for i in range(B):
_, info, _ = dynedc_compress_binary(spk_np[i], chunk_size=chunk_size)
per_layer[k]["ratios"].append(info["compression_ratio"])
per_layer[k]["modes"].append(info.get("mode", "unknown"))
summary = {"chunk_size": chunk_size, "tensor": "cadlif_spike_output"}
all_ratios = []
for k in layer_keys:
ratios = per_layer[k]["ratios"]
modes = per_layer[k]["modes"]
all_ratios.extend(ratios)
summary[k] = {
"mean_compression_ratio": float(np.mean(ratios)),
"mean_space_saving_pct": float((1 - np.mean(ratios)) * 100),
"min_ratio": float(np.min(ratios)),
"max_ratio": float(np.max(ratios)),
"n_samples": len(ratios),
"mode_distribution": dict(Counter(modes)),
}
summary["overall"] = {
"mean_compression_ratio": float(np.mean(all_ratios)),
"mean_space_saving_pct": float((1 - np.mean(all_ratios)) * 100),
"min_ratio": float(np.min(all_ratios)),
"max_ratio": float(np.max(all_ratios)),
"n_samples": len(all_ratios),
}
return summary
# =============================================================================
# DyNED Encoder Layer
# =============================================================================
class DyNEDEncoderLayer(nn.Module):
"""DyNED sigma-delta quantization as a differentiable PyTorch layer."""
def __init__(self, levels=256):
super().__init__()
self.levels = levels
def forward(self, x):
if not self.training:
return self._sigma_delta(x)
quantized = self._sigma_delta(x)
return x + (quantized - x).detach()
@torch.autocast("cuda", enabled=False)
def _sigma_delta(self, x):
x = x.float()
batch_size, n_features = x.shape
x_min = x.min(dim=1, keepdim=True).values
x_max = x.max(dim=1, keepdim=True).values
x_range = (x_max - x_min).clamp(min=1e-8)
x_norm = (x - x_min) / x_range
quantized = torch.zeros_like(x_norm)
error = torch.zeros(batch_size, 1, device=x.device, dtype=torch.float32)
levels_m1 = self.levels - 1
for i in range(n_features):
sample_with_error = x_norm[:, i : i + 1] + error
q = torch.round(sample_with_error * levels_m1).clamp(0, levels_m1) / levels_m1
quantized[:, i : i + 1] = q
error = sample_with_error - q
return quantized * x_range + x_min
# =============================================================================
# Surrogate Gradient
# =============================================================================
class ATanSurrogate(torch.autograd.Function):
@staticmethod
def forward(ctx, x, alpha=5.0):
ctx.save_for_backward(x)
ctx.alpha = alpha
return (x > 0).float()
@staticmethod
def backward(ctx, grad_output):
(x,) = ctx.saved_tensors
alpha = ctx.alpha
grad = alpha / (2.0 * (1.0 + (math.pi / 2.0 * alpha * x) ** 2))
return grad_output * grad, None
def spike_fn(x):
return ATanSurrogate.apply(x)
# =============================================================================
# cAdLIF Neuron
# =============================================================================
ALPHA_MIN = math.exp(-1.0 / 2.0) # 0.6065
ALPHA_MAX = math.exp(-1.0 / 25.0) # 0.9608
BETA_MIN = math.exp(-1.0 / 30.0) # 0.9672
BETA_MAX = math.exp(-1.0 / 120.0) # 0.9917
THRESHOLD = 0.5
class cAdLIFNeuron(nn.Module):
"""Constrained Adaptive LIF neuron (Deckers et al. 2024)."""
def __init__(self, size):
super().__init__()
self.size = size
self.alpha_raw = nn.Parameter(torch.empty(size).uniform_(ALPHA_MIN, ALPHA_MAX))
self.beta_raw = nn.Parameter(torch.empty(size).uniform_(BETA_MIN, BETA_MAX))
self.a_raw = nn.Parameter(torch.empty(size).uniform_(0.0, 0.5))
self.b_raw = nn.Parameter(torch.empty(size).uniform_(0.0, 1.0))
def _constrain(self):
alpha = self.alpha_raw.clamp(ALPHA_MIN, ALPHA_MAX)
beta = self.beta_raw.clamp(BETA_MIN, BETA_MAX)
a = self.a_raw.clamp(0.0, 1.0)
b = self.b_raw.clamp(0.0, 2.0)
return alpha, beta, a, b
def forward(self, I_t, u_prev, w_prev, s_prev):
alpha, beta, a, b = self._constrain()
w = beta * w_prev + (1.0 - beta) * a * u_prev + b * s_prev
u = alpha * u_prev + (1.0 - alpha) * (I_t - w_prev) - THRESHOLD * s_prev
s = spike_fn(u - THRESHOLD)
return s, u, w
# =============================================================================
# Learnable Delays
# =============================================================================
def apply_delays(h_seq, delays, max_delay):
T, B, N = h_seq.shape
device = h_seq.device
pad = torch.zeros(max_delay, B, N, device=device, dtype=h_seq.dtype)
h_padded = torch.cat([pad, h_seq], dim=0)
T_pad = T + max_delay
delays_clamped = delays.clamp(0.0, float(max_delay))
d_floor = delays_clamped.detach().long()
d_ceil = (d_floor + 1).clamp(max=max_delay)
frac = delays_clamped - d_floor.float()
t_range = torch.arange(T, device=device)
idx_floor = (t_range.unsqueeze(1) + max_delay - d_floor.unsqueeze(0)).clamp(0, T_pad - 1)
idx_ceil = (t_range.unsqueeze(1) + max_delay - d_ceil.unsqueeze(0)).clamp(0, T_pad - 1)
idx_f = idx_floor.unsqueeze(1).expand(T, B, N)
idx_c = idx_ceil.unsqueeze(1).expand(T, B, N)
val_floor = torch.gather(h_padded, 0, idx_f)
val_ceil = torch.gather(h_padded, 0, idx_c)
frac_exp = frac.view(1, 1, N)
return (1.0 - frac_exp) * val_floor + frac_exp * val_ceil
# =============================================================================
# TET Loss (Temporal Efficient Training, Deng et al. ICLR 2022)
# =============================================================================
def tet_loss(m_out, targets, label_smoothing=0.1, lambda_tet=1e-3):
"""TET loss: per-timestep CE averaged + temporal consistency MSE.
Standard SNN loss applies CE to the time-averaged output.
TET applies CE at each timestep independently, then averages.
This gives stronger gradients to earlier timesteps.
Args:
m_out: [T, B, num_outputs] - per-timestep readout membrane potentials
targets: [B] - class labels
label_smoothing: for cross-entropy
lambda_tet: weight for temporal consistency regularization
Returns:
loss: scalar
"""
T = m_out.shape[0]
# Per-timestep cross-entropy, averaged over time
ce_loss = sum(F.cross_entropy(m_out[t], targets, label_smoothing=label_smoothing) for t in range(T)) / T
# Temporal consistency: each timestep should predict similarly
m_mean = m_out.mean(dim=0, keepdim=True) # [1, B, C]
mse_reg = F.mse_loss(m_out, m_mean.expand_as(m_out))
return ce_loss + lambda_tet * mse_reg
# =============================================================================
# VGGSNN Conv + DyNED (temporal) + cAdLIF SNN
# =============================================================================
class VGGDyNEDcAdLIFSNN(nn.Module):
"""VGGSNN conv frontend + temporal DyNED encoding + cAdLIF SNN for CIFAR10-DVS.
VGGSNN consistently outperforms ResNet on CIFAR10-DVS in the SNN literature
(83-84% vs 74-78%). Deeper sequential conv without skip connections.
Architecture:
Input [B, T=16, 2, 32, 32] - temporal DVS frames (dual polarity)
-> reshape [B*T, 2, 32, 32]
-> VGG conv (2->64->128->256->256->512->512->512) with AvgPool
-> AdaptiveAvgPool(4x4) -> flatten -> 8192
-> proj_down(8192->dyned_dim) -> LayerNorm
-> reshape [B*dyned_dim, T] -> DyNED (sigma-delta across time)
-> reshape [B*T, dyned_dim] -> proj_up(dyned_dim->hidden) -> LayerNorm
-> [T, B, hidden]
-> cAdLIF layer 1 + delay + dropout
-> cAdLIF layer 2 + delay + dropout
-> cAdLIF layer 3 + delay + dropout
-> Linear readout + softmax accumulation (or TET per-timestep)
"""
def __init__(
self,
hidden_size=512,
num_outputs=10,
dyned_levels=256,
dyned_dim=512,
dropout=0.25,
conv_dropout=0.1,
max_delay=5,
dynedc_in_forward=False,
dynedc_chunk_size=4,
):
super().__init__()
self.hidden_size = hidden_size
self.num_outputs = num_outputs
self.max_delay = max_delay
self.dyned_dim = dyned_dim
self.dyned_levels = dyned_levels
self.dynedc_in_forward = dynedc_in_forward
self.dynedc_chunk_size = dynedc_chunk_size
# NB: DyNEDc is measured in a single post-training pass over the test
# set (see measure_compression_stats). The model is identical to v8
# DyNED - DyNEDc is informational and out of the data path during
# training because it is lossless and pre-determined by the model's
# learned DyNED outputs at the end of training.
# VGG-style conv backbone
self.features = nn.Sequential(
# Block 1: 2 -> 64 (32x32)
nn.Conv2d(_N_CHANNELS, 64, 3, padding=1, bias=False),
nn.BatchNorm2d(64),
nn.ReLU(),
# Block 2: 64 -> 128, pool -> 16x16
nn.Conv2d(64, 128, 3, padding=1, bias=False),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.AvgPool2d(2),
nn.Dropout2d(conv_dropout),
# Block 3: 128 -> 256 (16x16)
nn.Conv2d(128, 256, 3, padding=1, bias=False),
nn.BatchNorm2d(256),
nn.ReLU(),
# Block 4: 256 -> 256, pool -> 8x8
nn.Conv2d(256, 256, 3, padding=1, bias=False),
nn.BatchNorm2d(256),
nn.ReLU(),
nn.AvgPool2d(2),
nn.Dropout2d(conv_dropout),
# Block 5: 256 -> 512 (8x8)
nn.Conv2d(256, 512, 3, padding=1, bias=False),
nn.BatchNorm2d(512),
nn.ReLU(),
# Block 6: 512 -> 512, pool -> 4x4
nn.Conv2d(512, 512, 3, padding=1, bias=False),
nn.BatchNorm2d(512),
nn.ReLU(),
nn.AvgPool2d(2),
nn.Dropout2d(conv_dropout),
# Block 7: 512 -> 512 (4x4)
nn.Conv2d(512, 512, 3, padding=1, bias=False),
nn.BatchNorm2d(512),
nn.ReLU(),
)
conv_out_size = 512 * 4 * 4 # 8192
# DyNED path. The DyNED encoder produces continuous level-quantised
# values in [0, dyned_levels-1], which are projected up to the cAdLIF
# input width. The cAdLIF stack itself emits the binary spike trains
# that DyNEDc compresses post-training.
self.proj_down = nn.Linear(conv_out_size, dyned_dim)
self.proj_down_norm = nn.LayerNorm(dyned_dim)
self.dyned_encoder = DyNEDEncoderLayer(levels=dyned_levels)
self.proj_up = nn.Linear(dyned_dim, hidden_size)
self.proj_up_norm = nn.LayerNorm(hidden_size)
# cAdLIF layers with delays
# Delays initialised uniformly in [0, max_delay] (Hammouamri 2024
# "Learning Delays in Spiking Neural Networks") so each neuron starts
# at a different value and gradients can move it in either direction.
# Zero-init causes every delay to sit at the lower clamp boundary and
# never drift, since clamp masks any negative gradient back to 0.
self.fc1 = nn.Linear(hidden_size, hidden_size)
self.bn1 = nn.BatchNorm1d(hidden_size, momentum=0.05)
self.cadlif1 = cAdLIFNeuron(hidden_size)
self.delay1 = nn.Parameter(torch.empty(hidden_size).uniform_(0.0, float(max_delay)))
self.drop1 = nn.Dropout(dropout)
self.fc2 = nn.Linear(hidden_size, hidden_size)
self.bn2 = nn.BatchNorm1d(hidden_size, momentum=0.05)
self.cadlif2 = cAdLIFNeuron(hidden_size)
self.delay2 = nn.Parameter(torch.empty(hidden_size).uniform_(0.0, float(max_delay)))
self.drop2 = nn.Dropout(dropout)
self.fc3 = nn.Linear(hidden_size, hidden_size)
self.bn3 = nn.BatchNorm1d(hidden_size, momentum=0.05)
self.cadlif3 = cAdLIFNeuron(hidden_size)
self.delay3 = nn.Parameter(torch.empty(hidden_size).uniform_(0.0, float(max_delay)))
self.drop3 = nn.Dropout(dropout)
# Readout
self.fc_out = nn.Linear(hidden_size, num_outputs)
self.alpha_out_raw = nn.Parameter(torch.empty(num_outputs).uniform_(ALPHA_MIN, ALPHA_MAX))
# Init
nn.init.kaiming_uniform_(self.fc1.weight)
nn.init.kaiming_uniform_(self.fc2.weight)
nn.init.kaiming_uniform_(self.fc3.weight)
nn.init.kaiming_uniform_(self.fc_out.weight)
def _get_alpha_out(self):
return self.alpha_out_raw.clamp(ALPHA_MIN, ALPHA_MAX)
def _encode(self, x):
"""VGG conv + DyNED temporal encoding + in-forward ON/OFF binarisation.
Args:
x: [B, T, 2, 32, 32]
Returns:
x_seq: [T, B, hidden_size]
"""
B, T = x.shape[0], x.shape[1]
# Process all frames through shared VGG conv
x_flat = x.reshape(B * T, _N_CHANNELS, _TARGET_H, _TARGET_W) # [B*T, 2, 32, 32]
h = self.features(x_flat) # [B*T, 512, 4, 4]
# Project conv features to DyNED dim
compact = self.proj_down_norm(self.proj_down(h.flatten(1))) # [B*T, dyned_dim]
# DyNED across time: sigma-delta on temporal evolution of each feature
compact_3d = compact.reshape(B, T, self.dyned_dim)
compact_time = compact_3d.permute(0, 2, 1).reshape(B * self.dyned_dim, T) # [B*dyned_dim, T]
encoded_time = self.dyned_encoder(compact_time) # [B*dyned_dim, T]
encoded_3d = encoded_time.reshape(B, self.dyned_dim, T) # [B, dyned_dim, T]
# Continuous DyNED-quantised values feed straight into the cAdLIF stack.
# No in-forward ON/OFF binarisation: the cAdLIF layers themselves emit
# native binary spike trains, and DyNEDc compresses *those* outputs
# post-training (see `measure_compression_stats`).
encoded_flat = encoded_3d.permute(0, 2, 1).reshape(B * T, self.dyned_dim)
features = self.proj_up_norm(self.proj_up(encoded_flat)) # [B*T, hidden]
return features.reshape(B, T, self.hidden_size).permute(1, 0, 2) # [T, B, hidden]
def forward(self, x, return_temporal=False):
"""
Args:
x: [B, T, 2, 32, 32] - temporal DVS frames (dual polarity)
return_temporal: if True, also return per-timestep membrane outputs for TET loss
Returns:
output: [B, num_outputs] - accumulated softmax votes
m_out: [T, B, num_outputs] - (only if return_temporal=True)
"""
B, T = x.shape[0], x.shape[1]
device = x.device
x_seq = self._encode(x) # [T, B, hidden]
# Layer 1
h1 = self.fc1(x_seq)
h1 = self.bn1(h1.reshape(-1, self.hidden_size)).reshape(T, B, self.hidden_size)
h1 = apply_delays(h1, self.delay1.clamp(0, self.max_delay), self.max_delay)
s1_list = []
u1 = torch.zeros(B, self.hidden_size, device=device)
w1 = torch.zeros(B, self.hidden_size, device=device)
s1 = torch.zeros(B, self.hidden_size, device=device)
for t in range(T):
s1, u1, w1 = self.cadlif1(h1[t], u1, w1, s1)
s1_list.append(s1)
s1_seq = self.drop1(torch.stack(s1_list))
# Layer 2
h2 = self.fc2(s1_seq)
h2 = self.bn2(h2.reshape(-1, self.hidden_size)).reshape(T, B, self.hidden_size)
h2 = apply_delays(h2, self.delay2.clamp(0, self.max_delay), self.max_delay)
s2_list = []
u2 = torch.zeros(B, self.hidden_size, device=device)
w2 = torch.zeros(B, self.hidden_size, device=device)
s2 = torch.zeros(B, self.hidden_size, device=device)
for t in range(T):
s2, u2, w2 = self.cadlif2(h2[t], u2, w2, s2)
s2_list.append(s2)
s2_seq = self.drop2(torch.stack(s2_list))
# Layer 3
h3 = self.fc3(s2_seq)
h3 = self.bn3(h3.reshape(-1, self.hidden_size)).reshape(T, B, self.hidden_size)
h3 = apply_delays(h3, self.delay3.clamp(0, self.max_delay), self.max_delay)
s3_list = []
u3 = torch.zeros(B, self.hidden_size, device=device)
w3 = torch.zeros(B, self.hidden_size, device=device)
s3 = torch.zeros(B, self.hidden_size, device=device)
for t in range(T):
s3, u3, w3 = self.cadlif3(h3[t], u3, w3, s3)
s3_list.append(s3)
s3_seq = self.drop3(torch.stack(s3_list))
# Readout: LIF with infinite threshold
alpha_out = self._get_alpha_out()
u_out = torch.zeros(B, self.num_outputs, device=device)
m_out_list = []
for t in range(T):
cur = self.fc_out(s3_seq[t])
u_out = alpha_out * u_out + (1.0 - alpha_out) * cur
m_out_list.append(u_out)
m_out = torch.stack(m_out_list) # [T, B, num_outputs]
output = torch.sum(F.softmax(m_out, dim=2), dim=0)
if return_temporal:
return output, m_out
return output
def diagnostic_forward(self, x):
B, T = x.shape[0], x.shape[1]
device = x.device
x_seq = self._encode(x)
layer_data = {"spk1": [], "spk2": [], "spk3": [], "mem_out": []}
mem1_list = []
mem2_list = []
mem3_list = []
h1 = self.fc1(x_seq)
h1 = self.bn1(h1.reshape(-1, self.hidden_size)).reshape(T, B, self.hidden_size)
h1 = apply_delays(h1, self.delay1.clamp(0, self.max_delay), self.max_delay)
u1 = w1 = s1 = torch.zeros(B, self.hidden_size, device=device)
for t in range(T):
s1, u1, w1 = self.cadlif1(h1[t], u1, w1, s1)
layer_data["spk1"].append(s1)
mem1_list.append(u1)
s1_seq = torch.stack(layer_data["spk1"])
h2 = self.fc2(s1_seq)
h2 = self.bn2(h2.reshape(-1, self.hidden_size)).reshape(T, B, self.hidden_size)
h2 = apply_delays(h2, self.delay2.clamp(0, self.max_delay), self.max_delay)
u2 = w2 = s2 = torch.zeros(B, self.hidden_size, device=device)
for t in range(T):
s2, u2, w2 = self.cadlif2(h2[t], u2, w2, s2)
layer_data["spk2"].append(s2)
mem2_list.append(u2)
s2_seq = torch.stack(layer_data["spk2"])
h3 = self.fc3(s2_seq)
h3 = self.bn3(h3.reshape(-1, self.hidden_size)).reshape(T, B, self.hidden_size)
h3 = apply_delays(h3, self.delay3.clamp(0, self.max_delay), self.max_delay)
u3 = w3 = s3 = torch.zeros(B, self.hidden_size, device=device)
alpha_out = self._get_alpha_out()
u_out = torch.zeros(B, self.num_outputs, device=device)
for t in range(T):
s3, u3, w3 = self.cadlif3(h3[t], u3, w3, s3)
layer_data["spk3"].append(s3)
mem3_list.append(u3)
cur = self.fc_out(s3)
u_out = alpha_out * u_out + (1.0 - alpha_out) * cur
layer_data["mem_out"].append(u_out)
for key in layer_data:
layer_data[key] = torch.stack(layer_data[key])
layer_data["mem1"] = torch.stack(mem1_list)
layer_data["mem2"] = torch.stack(mem2_list)
layer_data["mem3"] = torch.stack(mem3_list)
return layer_data
# =============================================================================
# Data Setup
# =============================================================================
def setup_training(batch_size=128, workers=4, train_split=0.9, max_samples=None):
assets_dir = os.path.join("..", "assets")
os.makedirs(assets_dir, exist_ok=True)
print("Loading CIFAR10-DVS dataset...")
full_dataset = CIFAR10DVSFrameDataset(save_to=assets_dir, augment=False, cache=True)
n_total = len(full_dataset)
# Optional stratified-by-class subsetting for quick smoke tests.
if max_samples is not None and max_samples < n_total:
if hasattr(full_dataset, "_labels") and full_dataset._labels is not None:
labels_np = full_dataset._labels.numpy()
n_classes = len(CIFAR10DVS_CLASSES)
per_class = max(1, max_samples // n_classes)
rng = np.random.default_rng(42)
picked = []
for c in range(n_classes):
cls_idx = np.where(labels_np == c)[0]
if len(cls_idx) == 0:
continue
take = min(per_class, len(cls_idx))
picked.append(rng.choice(cls_idx, size=take, replace=False))
subset_idx = np.concatenate(picked) if picked else np.arange(min(max_samples, n_total))
rng.shuffle(subset_idx)
subset_idx = subset_idx[:max_samples]
full_dataset = torch.utils.data.Subset(full_dataset, subset_idx.tolist())
n_total = len(full_dataset)
print(f"** Subset mode: stratified subset of {n_total} samples")
else:
full_dataset = torch.utils.data.Subset(
full_dataset,
list(range(min(max_samples, n_total))),
)
n_total = len(full_dataset)
print(f"** Subset mode: first {n_total} samples")
n_train = int(n_total * train_split)
n_test = n_total - n_train
print(f"Dataset: {n_total} total - {n_train} train, {n_test} test")
generator = torch.Generator().manual_seed(42)
train_dataset, test_dataset = random_split(full_dataset, [n_train, n_test], generator=generator)
train_dataset_aug = _AugmentedSubset(train_dataset)
dl_kwargs = dict(pin_memory=True, persistent_workers=False)
if workers > 0:
dl_kwargs["prefetch_factor"] = 4
train_loader = DataLoader(
train_dataset_aug,
batch_size=batch_size,
shuffle=True,
num_workers=workers,
**dl_kwargs,
)
test_loader = DataLoader(
test_dataset,
batch_size=batch_size * 2,
shuffle=False,
num_workers=workers,
**dl_kwargs,
)
return train_loader, test_loader
# =============================================================================
# Training Loop (with TET loss)
# =============================================================================
def train_network(
net,
train_loader,
test_loader,
num_epochs=300,
device="cuda",
lr=0.1,
lr_delay=0.01,
weight_decay=5e-4,
lambda_tet=1e-3,
tet_start_epoch=50,
trial=None,
):
delay_params = [net.delay1, net.delay2, net.delay3]
# cAdLIF cell + readout dynamics parameters: physically meaningful, must NOT
# be weight-decayed toward zero. Decaying `a_raw`/`b_raw` collapses cAdLIF
# into a vanilla LIF and drops accuracy by ~15%.
neuron_params = [
net.cadlif1.alpha_raw,
net.cadlif1.beta_raw,
net.cadlif1.a_raw,
net.cadlif1.b_raw,
net.cadlif2.alpha_raw,
net.cadlif2.beta_raw,
net.cadlif2.a_raw,
net.cadlif2.b_raw,
net.cadlif3.alpha_raw,
net.cadlif3.beta_raw,
net.cadlif3.a_raw,
net.cadlif3.b_raw,
net.alpha_out_raw,
]
no_decay_ids = {id(p) for p in delay_params + neuron_params}
weight_params = [p for p in net.parameters() if id(p) not in no_decay_ids]
optimizer = torch.optim.SGD(
[
{"params": weight_params, "lr": lr, "weight_decay": weight_decay, "momentum": 0.9},
{"params": neuron_params, "lr": lr, "weight_decay": 0.0, "momentum": 0.9},
{"params": delay_params, "lr": lr_delay, "weight_decay": 0.0, "momentum": 0.9},
]
)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer,
T_max=num_epochs,
eta_min=1e-6,
)
scaler = torch.amp.GradScaler("cuda") if device == "cuda" else None
best_acc = 0
metrics = {
"epoch": [],
"train_loss": [],
"test_accuracy": [],
"learning_rate": [],
"epoch_time": [],
"layer_firing_rates": [],
"per_class_accuracy": [],
}
print(
f"Training on {device} | Batch size: {train_loader.batch_size} | "
f"SGD(momentum=0.9) | TET \u03bb={lambda_tet} starts at epoch {tet_start_epoch}"
)
best_model_state = None
for epoch in range(num_epochs):
epoch_start = time.time()
net.train()
epoch_loss = 0.0
num_batches = 0
running_loss = 0.0
use_tet = epoch >= tet_start_epoch
for i, (data, targets) in enumerate(train_loader):
data = data.to(device, non_blocking=True)
targets = targets.to(device, non_blocking=True)
optimizer.zero_grad()
if scaler is not None:
with torch.amp.autocast("cuda"):
output, m_out = net(data, return_temporal=True)
if use_tet:
loss = tet_loss(m_out, targets, label_smoothing=0.1, lambda_tet=lambda_tet)
else:
# Standard CE on time-averaged membrane potentials (logits)
loss = F.cross_entropy(m_out.mean(dim=0), targets, label_smoothing=0.1)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
else:
output, m_out = net(data, return_temporal=True)
if use_tet:
loss = tet_loss(m_out, targets, label_smoothing=0.1, lambda_tet=lambda_tet)
else:
loss = F.cross_entropy(m_out.mean(dim=0), targets, label_smoothing=0.1)
loss.backward()
torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)
optimizer.step()
batch_loss = loss.item()
running_loss += batch_loss
epoch_loss += batch_loss
num_batches += 1
if i % 50 == 49:
avg_loss = running_loss / 50
print(f"Epoch {epoch + 1}, Batch {i + 1}: Loss = {avg_loss:.4f}", end="")
if device == "cuda":
print(f" | GPU: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
else:
print()
running_loss = 0.0
epoch_time = time.time() - epoch_start
avg_epoch_loss = epoch_loss / num_batches
current_lr = optimizer.param_groups[0]["lr"]
# Firing rates diagnostic
with torch.no_grad():
net.eval()
diag_data = next(iter(test_loader))[0][:16].to(device)
layer_data = net.diagnostic_forward(diag_data)
firing_rates = {
"layer1": layer_data["spk1"].mean().item(),
"layer2": layer_data["spk2"].mean().item(),
"layer3": layer_data["spk3"].mean().item(),
}
eval_result = evaluate(net, test_loader, device)
test_acc = eval_result["accuracy"]
scheduler.step()
metrics["epoch"].append(epoch + 1)
metrics["train_loss"].append(avg_epoch_loss)
metrics["test_accuracy"].append(test_acc)
metrics["learning_rate"].append(current_lr)
metrics["epoch_time"].append(epoch_time)
metrics["layer_firing_rates"].append(firing_rates)
metrics["per_class_accuracy"].append(eval_result["per_class_accuracy"])
alpha1 = net.cadlif1.alpha_raw.clamp(ALPHA_MIN, ALPHA_MAX).mean().item()
alpha2 = net.cadlif2.alpha_raw.clamp(ALPHA_MIN, ALPHA_MAX).mean().item()
alpha3 = net.cadlif3.alpha_raw.clamp(ALPHA_MIN, ALPHA_MAX).mean().item()
d1_mean = net.delay1.clamp(0, net.max_delay).mean().item()
d2_mean = net.delay2.clamp(0, net.max_delay).mean().item()
d3_mean = net.delay3.clamp(0, net.max_delay).mean().item()
phase = "TET" if use_tet else "CE"
print(
f"Epoch {epoch + 1} [{phase}]: Test Acc = {test_acc:.4f} | Loss = {avg_epoch_loss:.4f} | "
f"LR = {current_lr:.6f} | \u03b1 = [{alpha1:.3f}, {alpha2:.3f}, {alpha3:.3f}] | "
f"Delay = [{d1_mean:.1f}, {d2_mean:.1f}, {d3_mean:.1f}] | "
f"FR = [{firing_rates['layer1']:.3f}, {firing_rates['layer2']:.3f}, "
f"{firing_rates['layer3']:.3f}] | {epoch_time:.1f}s"
)
if device == "cuda":
torch.cuda.empty_cache()
if test_acc > best_acc:
best_acc = test_acc
best_model_state = {k: v.cpu().clone() for k, v in net.state_dict().items()}
torch.save(
{
"epoch": epoch,
"model_state_dict": net.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"accuracy": test_acc,
},
os.path.join(OUTPUT_DIR, "best_model.pth"),
)
if trial is not None:
import optuna
trial.report(test_acc, epoch)
if trial.should_prune():
raise optuna.TrialPruned()
if best_model_state is not None:
net.load_state_dict(best_model_state)
return metrics
# =============================================================================
# Evaluation
# =============================================================================
def evaluate(net, test_loader, device, collect_representations=False):
net.eval()
correct = 0
total = 0
all_predictions = []
all_targets = []
all_representations = []
per_class_correct = np.zeros(10)
per_class_total = np.zeros(10)
with torch.no_grad():
for data, targets in test_loader:
data = data.to(device, non_blocking=True)
targets = targets.to(device, non_blocking=True)
output = net(data)
predicted_classes = output.argmax(dim=1)
correct += (predicted_classes == targets).sum().item()
total += targets.size(0)
all_predictions.extend(predicted_classes.cpu().numpy())
all_targets.extend(targets.cpu().numpy())
for cls in range(10):
mask = targets == cls
per_class_correct[cls] += (predicted_classes[mask] == targets[mask]).sum().item()
per_class_total[cls] += mask.sum().item()
if collect_representations:
all_representations.append(output.cpu())
accuracy = correct / total if total > 0 else 0.0
per_class_acc = per_class_correct / (per_class_total + 1e-8)
all_predictions = np.array(all_predictions)
all_targets = np.array(all_targets)
cm = np.zeros((10, 10), dtype=np.int64)
for pred_cls, true_cls in zip(all_predictions, all_targets):
cm[true_cls][pred_cls] += 1
result = {
"accuracy": accuracy,
"per_class_accuracy": {CIFAR10DVS_CLASSES[i]: float(per_class_acc[i]) for i in range(10)},
"confusion_matrix": cm,
"predictions": all_predictions,
"targets": all_targets,
}
if collect_representations and all_representations:
result["representations"] = torch.cat(all_representations, dim=0).numpy()
return result
# =============================================================================
# Visualizations
# =============================================================================
def analyze_training_metrics(metrics):
losses = np.array(metrics["train_loss"])
accuracies = np.array(metrics["test_accuracy"])
epochs = np.array(metrics["epoch"])
header_parts = ["epoch", "train_loss", "test_accuracy", "learning_rate", "epoch_time"]
fr_keys = sorted(metrics["layer_firing_rates"][0].keys()) if metrics["layer_firing_rates"] else []
header_parts.extend(f"firing_rate_{k}" for k in fr_keys)
header_parts.extend(f"acc_{cls}" for cls in CIFAR10DVS_CLASSES)
header = ",".join(header_parts)
rows = []
for i in range(len(metrics["epoch"])):
row = [
metrics["epoch"][i],
metrics["train_loss"][i],
metrics["test_accuracy"][i],
metrics["learning_rate"][i],
metrics["epoch_time"][i],
]
for k in fr_keys:
row.append(metrics["layer_firing_rates"][i].get(k, 0))
pca = metrics["per_class_accuracy"][i]
for cls in CIFAR10DVS_CLASSES:
row.append(pca.get(cls, 0))
rows.append(row)
filepath = os.path.join(OUTPUT_DIR, "training_metrics.csv")
np.savetxt(filepath, np.array(rows), delimiter=",", header=header, comments="")
best_epoch = np.argmax(accuracies)
print(f"\nBest epoch: {best_epoch + 1} - test accuracy: {accuracies[best_epoch]:.4f}")
fig, ax1 = plt.subplots(figsize=(12, 6))
ax1.set_xlabel("Epoch")
ax1.set_ylabel("Loss", color="tab:blue")
ax1.plot(epochs, losses, color="tab:blue", label="Training Loss (TET)", alpha=0.8)
ax1.tick_params(axis="y", labelcolor="tab:blue")
ax2 = ax1.twinx()
ax2.set_ylabel("Accuracy (%)", color="tab:red")
ax2.plot(epochs, accuracies * 100, color="tab:red", label="Test Accuracy", alpha=0.8)
ax2.tick_params(axis="y", labelcolor="tab:red")
plt.title("CIFAR10-DVS VGGSNN+DyNED+cAdLIF+TET SNN with DyNEDc Training Progress")
lines1, labels1 = ax1.get_legend_handles_labels()
lines2, labels2 = ax2.get_legend_handles_labels()
ax2.legend(lines1 + lines2, labels1 + labels2, loc="upper right")
plt.savefig(os.path.join(OUTPUT_DIR, "training_progress.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR,
"training_progress",
epochs=epochs,
losses=losses,
accuracies=accuracies,
)
if metrics["layer_firing_rates"]:
fr_history = metrics["layer_firing_rates"]
keys = sorted(fr_history[0].keys())
plt.figure(figsize=(12, 6))
fr_per_key = {}
for key in keys:
rates = [fr[key] for fr in fr_history]
fr_per_key[key] = np.asarray(rates)
plt.plot(epochs, rates, label=key, linewidth=1.5)
plt.xlabel("Epoch")
plt.ylabel("Mean Firing Rate")
plt.title("Per-Layer Firing Rates")
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "firing_rates.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR,
"firing_rates",
epochs=epochs,
keys=np.array(keys),
**{f"rate_{k}": fr_per_key[k] for k in keys},
)
plt.figure(figsize=(10, 4))
plt.plot(epochs, metrics["learning_rate"])
plt.xlabel("Epoch")
plt.ylabel("Learning Rate")
plt.title("Learning Rate Schedule")
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "lr_schedule.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR,
"lr_schedule",
epochs=epochs,
learning_rate=np.asarray(metrics["learning_rate"]),
)
def collect_diagnostic_batches(net, test_loader, device, max_samples=2000):
"""Run diagnostic_forward across multiple test batches, concatenating
per-key outputs along the batch dim. Used for per-class statistics."""
net.eval()
accumulated = {}
inputs = []
targets = []
n = 0
with torch.no_grad():
for data, tgt in test_loader:
data = data.to(device)
ld = net.diagnostic_forward(data)
for k, v in ld.items():
accumulated.setdefault(k, []).append(v.cpu())
inputs.append(data.cpu())
targets.append(tgt)
n += data.size(0)
if n >= max_samples:
break
out = {k: torch.cat(v, dim=1) for k, v in accumulated.items()} # cat over batch dim (T, B, ...)
return out, torch.cat(inputs, dim=0), torch.cat(targets, dim=0)
def visualize_network_activity(input_data, layer_data):
"""Four-panel inference snapshot for sample 0: input raster, hidden L3
raster, hidden L3 firing-rate distribution, output membrane potentials."""
spk_l3 = layer_data["spk3"] # [T, B, hidden]
mem_out = layer_data["mem_out"] # [T, B, num_classes]
hidden_w = spk_l3.shape[-1]
# input_data: [B, T, 2, H, W] -> sample 0 collapsed to [features, T] for raster
sample0 = input_data[0].detach().cpu().numpy() # [T, 2, H, W]
T_in = sample0.shape[0]
input_raster = sample0.reshape(T_in, -1).T # [features, T]
fig, axes = plt.subplots(2, 2, figsize=(15, 12))
axes[0, 0].imshow(input_raster, aspect="auto", origin="lower", cmap="binary")
axes[0, 0].set_title("Input Spike Train (Sample 0)")
axes[0, 0].set_xlabel("Time Frame")
axes[0, 0].set_ylabel("Spatial Feature (2x32x32 flat)")
axes[0, 1].imshow(spk_l3[:, 0].cpu().numpy(), aspect="auto", cmap="binary")
axes[0, 1].set_title(f"Hidden Layer 3 Spike Raster (Sample 0, {hidden_w})")
axes[0, 1].set_xlabel("Neuron Index")
axes[0, 1].set_ylabel("Time Step")
rates = spk_l3.mean(dim=0).cpu().numpy().flatten()
axes[1, 0].hist(rates, bins=50, color="steelblue")
axes[1, 0].set_title(f"Hidden Layer 3 Firing Rate Distribution ({hidden_w})")
axes[1, 0].set_xlabel("Firing Rate")
axes[1, 0].set_ylabel("Count")
axes[1, 1].plot(mem_out[:, 0].cpu().numpy())
axes[1, 1].set_title("Output Membrane Potentials per Class (Sample 0)")
axes[1, 1].set_xlabel("Time Step")
axes[1, 1].set_ylabel("Membrane Potential")
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "network_activity.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR,
"network_activity",
input_sample0=input_raster,
spk_l3_sample0=spk_l3[:, 0],
firing_rates=rates,
mem_out_sample0=mem_out[:, 0],
)
print("Saved network activity")
def visualize_layer_spike_rasters(layer_data):
"""Four-panel raster: spk1, spk2, spk3, mem_out (heatmap) for sample 0."""
fig, axes = plt.subplots(2, 2, figsize=(16, 12))
h1_w = layer_data["spk1"].shape[-1]
h2_w = layer_data["spk2"].shape[-1]
h3_w = layer_data["spk3"].shape[-1]
panels = [
("spk1", f"Layer 1 (Hidden, {h1_w})"),
("spk2", f"Layer 2 (Hidden, {h2_w})"),
("spk3", f"Layer 3 (Hidden, {h3_w})"),
("mem_out", "Output Layer (Membrane Potentials)"),
]
panel_arrs = {}
for ax, (key, title) in zip(axes.flat, panels):
d = layer_data[key][:, 0].cpu().numpy()
panel_arrs[key] = d
cmap = "binary" if key.startswith("spk") else "viridis"
ax.imshow(d, aspect="auto", cmap=cmap, interpolation="nearest")
ax.set_title(title)
ax.set_xlabel("Neuron Index")
ax.set_ylabel("Time Step")
if key.startswith("spk"):
ax.text(
0.02,
0.98,
f"Rate: {d.mean():.3f}",
transform=ax.transAxes,
va="top",
fontsize=9,
color="red",
bbox=dict(boxstyle="round", facecolor="white", alpha=0.8),
)
plt.suptitle("Per-Layer Activity (Sample 0)", fontsize=14)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "layer_spike_rasters.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR,
"layer_spike_rasters",
spk1=panel_arrs["spk1"],
spk2=panel_arrs["spk2"],
spk3=panel_arrs["spk3"],
mem_out=panel_arrs["mem_out"],
)
print("Saved layer spike rasters")
def visualize_membrane_distributions(layer_data):
"""Spike-count histograms for hidden layers (cAdLIF exposes spikes, not
membrane potentials internally) plus the output-layer membrane distribution.
Hidden-layer panels overlay the cAdLIF firing THRESHOLD as a vertical line."""
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
h1_w = layer_data["spk1"].shape[-1]
h2_w = layer_data["spk2"].shape[-1]
h3_w = layer_data["spk3"].shape[-1]
panels = [
("spk1", f"Layer 1 ({h1_w}) Spike Counts per Neuron", "spike"),
("spk2", f"Layer 2 ({h2_w}) Spike Counts per Neuron", "spike"),
("spk3", f"Layer 3 ({h3_w}) Spike Counts per Neuron", "spike"),
("mem_out", "Output Membrane Potentials", "membrane"),
]
panel_data = {}
for ax, (key, title, kind) in zip(axes.flat, panels):
d = layer_data[key].cpu().numpy()
if kind == "spike":
counts = d.sum(axis=0).flatten()
panel_data[key] = counts
ax.hist(counts, bins=50, density=True, color="steelblue", alpha=0.8)
ax.set_xlabel("Spike Count over Trial")
else:
vals = d.flatten()
panel_data[key] = vals
ax.hist(vals, bins=100, density=True, color="darkorange", alpha=0.8)
ax.axvline(THRESHOLD, color="red", linestyle="--", linewidth=1.0, label=f"THRESHOLD={THRESHOLD}")
ax.legend(loc="upper right", fontsize=8)
ax.set_xlabel("Membrane Potential")
ax.set_title(title)
ax.set_ylabel("Density")
ax.text(
0.02,
0.98,
f"Mean: {d.mean():.3f}\nStd: {d.std():.3f}",
transform=ax.transAxes,
va="top",
fontsize=9,
bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.8),
)
plt.suptitle("Activity Distributions", fontsize=14)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "membrane_distributions.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR,
"membrane_distributions",
spk1_counts=panel_data["spk1"],
spk2_counts=panel_data["spk2"],
spk3_counts=panel_data["spk3"],
mem_out_values=panel_data["mem_out"],
threshold=np.array(THRESHOLD),
)
print("Saved membrane distributions")
def visualize_per_class_spikes(layer_data, targets):
"""Per-class average spike patterns (output layer, binarised at THRESHOLD).
CIFAR10-DVS has 10 classes, plotted on a 2x5 grid."""
mem_out = layer_data["mem_out"] # [T, B, num_classes]
if isinstance(targets, torch.Tensor):
targets = targets.cpu()
class_labels = CIFAR10DVS_CLASSES if "CIFAR10DVS_CLASSES" in globals() else [str(i) for i in range(10)]
n_classes = len(class_labels)
rows, cols = 2, 5
fig, axes = plt.subplots(rows, cols, figsize=(3 * cols, 3 * rows))
T_dim = mem_out.shape[0]
C_dim = mem_out.shape[2]
per_class_arr = np.full((n_classes, T_dim, C_dim), np.nan, dtype=np.float32)
for cls_idx, ax in enumerate(axes.flat):
if cls_idx >= n_classes:
ax.axis("off")
continue
mask = targets == cls_idx
if mask.sum() == 0:
ax.set_title(class_labels[cls_idx], fontsize=8)
ax.text(0.5, 0.5, "no samples", transform=ax.transAxes, ha="center", va="center", fontsize=7)
continue
class_spikes = (mem_out[:, mask, :] > THRESHOLD).float().mean(dim=1).cpu().numpy()
per_class_arr[cls_idx] = class_spikes
ax.imshow(class_spikes, aspect="auto", cmap="binary", interpolation="nearest")
ax.set_title(class_labels[cls_idx], fontsize=8)
ax.tick_params(labelsize=6)
plt.suptitle("Per-Class Average Spike Patterns (Output Layer)", fontsize=14)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "per_class_spikes.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR,
"per_class_spikes",
per_class_arr=per_class_arr,
class_labels=np.array(class_labels),
grid=np.array([rows, cols]),
)
print("Saved per-class spikes")
def visualize_weight_distributions(net):
"""Trained weight histograms for the learned linear layers feeding each
cAdLIF block plus the readout."""
panels = [
("fc1", net.fc1.weight),
("fc2", net.fc2.weight),
("fc3", net.fc3.weight),
("fc_out", net.fc_out.weight),
]
n = len(panels)
fig, axes = plt.subplots(1, n, figsize=(5 * n, 5))
if n == 1:
axes = [axes]
weight_payload = {}
name_list = []
shape_list = []
for ax, (name, w) in zip(axes, panels):
vals = w.detach().cpu().numpy().flatten()
weight_payload[f"weight__{name}"] = vals
name_list.append(name)
shape_list.append(list(w.shape))
ax.hist(vals, bins=100, density=True, color="steelblue", alpha=0.8)
ax.set_title(f"{name} {tuple(w.shape)}")
ax.set_xlabel("Weight Value")
ax.set_ylabel("Density")
ax.text(
0.02,
0.98,
f"Mean: {vals.mean():.3f}\nStd: {vals.std():.3f}",
transform=ax.transAxes,
va="top",
fontsize=9,
bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.8),
)
plt.suptitle("Trained Weight Distributions", fontsize=14)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "weight_distributions.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR, "weight_distributions", names=np.array(name_list), shapes=np.array(shape_list), **weight_payload
)
print("Saved weight distributions")
def visualize_confusion_matrix_plot(cm):
fig, ax = plt.subplots(figsize=(10, 8))
cm_normalized = cm.astype(float) / (cm.sum(axis=1, keepdims=True) + 1e-8)
im = ax.imshow(cm_normalized, interpolation="nearest", cmap="Blues")
ax.figure.colorbar(im, ax=ax)
ax.set(
xticks=range(10),
yticks=range(10),
xticklabels=CIFAR10DVS_CLASSES,
yticklabels=CIFAR10DVS_CLASSES,
ylabel="True Label",
xlabel="Predicted Label",
title="Confusion Matrix (Normalized)",
)
plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "confusion_matrix.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR,
"confusion_matrix",
cm=cm,
cm_normalized=cm_normalized,
class_labels=np.array(CIFAR10DVS_CLASSES),
)
np.savetxt(
os.path.join(OUTPUT_DIR, "confusion_matrix.csv"),
cm,
delimiter=",",
fmt="%d",
header=",".join(CIFAR10DVS_CLASSES),
comments="",
)
print("Saved confusion matrix")
def visualize_tsne(representations, targets):
if not HAS_TSNE:
return
max_samples = 5000
if len(representations) > max_samples:
indices = np.random.choice(len(representations), max_samples, replace=False)
representations = representations[indices]
targets = targets[indices]
print("Computing t-SNE...")
tsne = TSNE(n_components=2, random_state=42, perplexity=30)
embedded = tsne.fit_transform(representations)
plt.figure(figsize=(12, 10))
scatter = plt.scatter(embedded[:, 0], embedded[:, 1], c=targets, cmap="tab10", alpha=0.6, s=5)
cbar = plt.colorbar(scatter, ticks=range(10))
cbar.set_ticklabels(CIFAR10DVS_CLASSES)
plt.title("t-SNE of Output Representations")
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "tsne.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR,
"tsne",
embedded=embedded,
targets=targets,
class_labels=np.array(CIFAR10DVS_CLASSES),
)
print("Saved t-SNE")
def visualize_per_class_accuracy(per_class_acc_dict):
labels = list(per_class_acc_dict.keys())
accs = [per_class_acc_dict[l] for l in labels]
sorted_pairs = sorted(zip(accs, labels), reverse=True)
accs_sorted = [p[0] for p in sorted_pairs]
labels_sorted = [p[1] for p in sorted_pairs]
plt.figure(figsize=(12, 6))
plt.bar(range(len(labels_sorted)), [a * 100 for a in accs_sorted], color="steelblue")
plt.xticks(range(len(labels_sorted)), labels_sorted, rotation=45, fontsize=9)
plt.ylabel("Accuracy (%)")
plt.title("Per-Class Accuracy on Test Set")
plt.ylim(0, 100)
plt.grid(True, axis="y", alpha=0.3)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "per_class_accuracy.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR,
"per_class_accuracy",
labels_sorted=np.array(labels_sorted),
accs_sorted=np.asarray(accs_sorted),
labels=np.array(labels),
accs=np.asarray(accs),
)
# =============================================================================
# Main
# =============================================================================
def main(json_path=None):
torch.manual_seed(42)
torch.cuda.manual_seed_all(42)
np.random.seed(42)
device_str = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device_str)
torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.set_float32_matmul_precision("high")
num_epochs = 300
tet_start_epoch = 50
if json_path:
import json
with open(json_path) as f:
params = json.load(f)
print(f"Loaded hyperparameters from: {json_path}")
lr = params["lr"]
lr_delay = params.get("lr_delay", 0.01)
weight_decay = params.get("weight_decay", 5e-4)
hidden_size = params["hidden_size"]
dyned_dim = params.get("dyned_dim", 512)
max_delay = params.get("max_delay", 5)
dropout = params.get("dropout", 0.25)
conv_dropout = params.get("conv_dropout", 0.1)
batch_size = params.get("batch_size", 128)
dyned_levels = params.get("dyned_levels", 256)
lambda_tet = params.get("lambda_tet", 1e-3)
tet_start_epoch = params.get("tet_start_epoch", 50)
dynedc_chunk_size = params.get("dynedc_chunk_size", 4)
else:
lr = 0.1
lr_delay = 0.01
weight_decay = 5e-4
hidden_size = 512
dyned_dim = 512
max_delay = 5
dropout = 0.25
conv_dropout = 0.1
batch_size = 128
# Lower dyned_levels: each surviving ON/OFF transition encodes a
# larger continuous-value change. With T=16, sparse-but-meaningful
# transitions seem to learn better than fine-grained ones.
dyned_levels = 64
lambda_tet = 1e-3
dynedc_chunk_size = 4
num_outputs = 10
workers = min(4, os.cpu_count() - 1) if os.cpu_count() > 1 else 0
# Quick-iteration overrides for smoke tests.
if _CLI_QUICK_EPOCHS is not None:
num_epochs = _CLI_QUICK_EPOCHS
tet_start_epoch = max(1, int(round(num_epochs * 50 / 300)))
print(f"** Quick mode: overriding num_epochs to {num_epochs} (TET starts at epoch {tet_start_epoch})")
max_samples = _CLI_SUBSET_SIZE
if max_samples is not None:
print(f"** Subset mode: limiting train+test split to first {max_samples} samples")
dynedc_in_forward = bool(_CLI_DYNEDC_IN_FORWARD)
if dynedc_in_forward:
print(
f"** DyNEDc in-forward enabled (chunk_size={dynedc_chunk_size}); "
f"compress+decompress every batch (lossless, slow)."
)
print("=" * 80)
print("VGGSNN+DyNED(2ch,T=16)+cAdLIF+TET for CIFAR10-DVS - v9 (DyNEDc post-training)")
print("=" * 80)
print(f"Device: {device_str.upper()}")
if torch.cuda.is_available():
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(
f"Training: DVS Frames [16\u00d72\u00d732\u00d732] \u2192 VGGSNN(2\u219264\u2192128\u2192256\u2192512) \u2192 DyNED(lvl={dyned_levels}) \u2192 cAdLIF SNN"
)
print(
f"Post-training: forward best model \u2192 cAdLIF spike outputs \u2192 DyNEDc(chunk={dynedc_chunk_size}) \u2192 stats JSON"
)
print(
f"Architecture: VGGSNN(7 layers) \u2192 DyNED({dyned_dim}\u00d7T=16) \u2192 {hidden_size} \u2192 {hidden_size} \u2192 {hidden_size} \u2192 {num_outputs}"
)
print(f"Timesteps: 16 | Delays: max={max_delay} | Dropout: {dropout}")
print(f"Batch: {batch_size} | LR: {lr} (weights) / {lr_delay} (delays) | SGD(momentum=0.9)")
print(
f"Loss: CE (epochs 1-{tet_start_epoch}) -> TET \u03bb={lambda_tet} (epochs {tet_start_epoch + 1}-{num_epochs})"
)
print(f"Augmentation: flip + affine + cutout(6\u00d76) + EventDrop(10%)")
print("=" * 80)
net = VGGDyNEDcAdLIFSNN(
hidden_size=hidden_size,
num_outputs=num_outputs,
dyned_levels=dyned_levels,
dyned_dim=dyned_dim,
dropout=dropout,
conv_dropout=conv_dropout,
max_delay=max_delay,
dynedc_in_forward=dynedc_in_forward,
dynedc_chunk_size=dynedc_chunk_size,
).to(device)
param_count = sum(p.numel() for p in net.parameters() if p.requires_grad)
print(f"Trainable parameters: {param_count:,}")
train_loader, test_loader = setup_training(
batch_size=batch_size,
workers=workers,
max_samples=max_samples,
)
try:
metrics = train_network(
net=net,
train_loader=train_loader,
test_loader=test_loader,
num_epochs=num_epochs,
device=device_str,
lr=lr,
lr_delay=lr_delay,
weight_decay=weight_decay,
lambda_tet=lambda_tet,
tet_start_epoch=tet_start_epoch,
)
print("\n" + "=" * 80)
print("POST-TRAINING ANALYSIS")
print("=" * 80)
analyze_training_metrics(metrics)
print("\nEvaluating best model on test set...")
test_eval = evaluate(net, test_loader, device, collect_representations=True)
print(f"Test accuracy: {test_eval['accuracy']:.4f}")
# Post-training DyNEDc compression measurement: one pass over the test
# set, extracting the model's DyNED output and compressing it. Because
# DyNEDc is lossless, this number reflects what the cache/transmission
# pipeline would achieve without affecting model accuracy.
print(f"\nMeasuring DyNEDc compression (chunk_size={dynedc_chunk_size})...")
stats = measure_compression_stats(
net,
test_loader,
device,
dyned_levels=dyned_levels,
chunk_size=dynedc_chunk_size,
)
print("DyNEDc compression statistics (post-training, test set, cAdLIF spike outputs):")
for layer_key in ("spk1", "spk2", "spk3", "overall"):
ls = stats[layer_key]
print(
f" {layer_key:<8}: ratio = {ls['mean_compression_ratio']:.4f} "
f"saving = {ls['mean_space_saving_pct']:.2f}% "
f"range = [{ls['min_ratio']:.4f}, {ls['max_ratio']:.4f}] "
f"n = {ls['n_samples']}"
)
import json as _json
with open(os.path.join(OUTPUT_DIR, "dynedc_compression_stats.json"), "w") as _fh:
_json.dump(stats, _fh, indent=2)
visualize_confusion_matrix_plot(test_eval["confusion_matrix"])
if "representations" in test_eval:
visualize_tsne(test_eval["representations"], test_eval["targets"])
visualize_per_class_accuracy(test_eval["per_class_accuracy"])
# Network dynamics figures from a single representative batch + an
# aggregated multi-batch view for per-class statistics.
net.eval()
with torch.no_grad():
sample_data, _ = next(iter(test_loader))
sample_data = sample_data.to(device)
sample_layer_data = net.diagnostic_forward(sample_data)
visualize_network_activity(sample_data, sample_layer_data)
visualize_layer_spike_rasters(sample_layer_data)
visualize_membrane_distributions(sample_layer_data)
agg_layer_data, _, agg_targets = collect_diagnostic_batches(
net,
test_loader,
device,
max_samples=2000,
)
visualize_per_class_spikes(agg_layer_data, agg_targets)
visualize_weight_distributions(net)
print("\nFinal Per-Class Accuracy (test set):")
for cls, acc in sorted(test_eval["per_class_accuracy"].items(), key=lambda x: x[1], reverse=True):
print(f" {cls:>12s}: {acc:.4f}")
print("\nLearned neuron parameters:")
for i, cadlif in enumerate([net.cadlif1, net.cadlif2, net.cadlif3], 1):
alpha, beta, a, b = cadlif._constrain()
print(f" Layer {i}: \u03b1={alpha.mean():.4f} \u03b2={beta.mean():.4f} a={a.mean():.4f} b={b.mean():.4f}")
print(f" Readout \u03b1: {net._get_alpha_out().mean():.4f}")
print("\nLearned delays:")
for i, delay in enumerate([net.delay1, net.delay2, net.delay3], 1):
d = delay.clamp(0, net.max_delay)
print(f" Layer {i}: mean={d.mean():.2f}, std={d.std():.2f}, range=[{d.min():.1f}, {d.max():.1f}]")
torch.save(
{
"model_state_dict": net.state_dict(),
"dyned_levels": dyned_levels,
"metrics": {
"train_losses": metrics["train_loss"],
"test_accuracies": metrics["test_accuracy"],
"test_accuracy": test_eval["accuracy"],
},
},
os.path.join(OUTPUT_DIR, "final_model.pth"),
)
print(f"\nDone! Best test accuracy: {max(metrics['test_accuracy']):.4f}")
print(f"Outputs: {OUTPUT_DIR}")
except KeyboardInterrupt:
print("Interrupted - saving...")
torch.save({"model_state_dict": net.state_dict()}, os.path.join(OUTPUT_DIR, "interrupted_model.pth"))
# =============================================================================
# Optuna
# =============================================================================
def optuna_objective(trial):
import optuna
device_str = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device_str)
lr = trial.suggest_float("lr", 0.01, 0.3, log=True)
lr_delay = trial.suggest_float("lr_delay", 1e-3, 0.1, log=True)
weight_decay = trial.suggest_float("weight_decay", 1e-5, 1e-3, log=True)
hidden_size = trial.suggest_categorical("hidden_size", [256, 512, 1024])
dyned_dim = trial.suggest_categorical("dyned_dim", [256, 512, 1024])
max_delay = trial.suggest_categorical("max_delay", [2, 3, 5, 7])
dropout = trial.suggest_float("dropout", 0.1, 0.4)
conv_dropout = trial.suggest_float("conv_dropout", 0.05, 0.2)
batch_size = trial.suggest_categorical("batch_size", [64, 128, 256])
dyned_levels = trial.suggest_categorical("dyned_levels", [16, 32, 64, 128, 256])
lambda_tet = trial.suggest_float("lambda_tet", 1e-4, 1e-1, log=True)
# NB: chunk_size only affects compression ratio, not model accuracy (DyNEDc
# is lossless and measured post-training). Kept in Optuna so best_params
# records the chunk_size used for the run, but accuracy is invariant to it.
_dynedc_chunk_size = trial.suggest_categorical("dynedc_chunk_size", [2, 4, 8])
# Optuna uses shorter runs - TET from epoch ~8 (proportional to 50/300)
num_epochs = 50
tet_start_epoch = max(1, int(round(num_epochs * 50 / 300)))
if _CLI_QUICK_EPOCHS is not None:
num_epochs = _CLI_QUICK_EPOCHS
tet_start_epoch = max(1, int(round(num_epochs * 50 / 300)))
workers = 0
net = VGGDyNEDcAdLIFSNN(
hidden_size=hidden_size,
num_outputs=10,
dyned_levels=dyned_levels,
dyned_dim=dyned_dim,
dropout=dropout,
conv_dropout=conv_dropout,
max_delay=max_delay,
dynedc_in_forward=bool(_CLI_DYNEDC_IN_FORWARD),
dynedc_chunk_size=_dynedc_chunk_size,
).to(device)
train_loader, test_loader = setup_training(
batch_size=batch_size,
workers=workers,
max_samples=_CLI_SUBSET_SIZE,
)
try:
metrics = train_network(
net=net,
train_loader=train_loader,
test_loader=test_loader,
num_epochs=num_epochs,
device=device_str,
lr=lr,
lr_delay=lr_delay,
weight_decay=weight_decay,
lambda_tet=lambda_tet,
tet_start_epoch=tet_start_epoch,
trial=trial,
)
except optuna.TrialPruned:
raise
finally:
del net, train_loader, test_loader
gc.collect()
if device_str == "cuda":
torch.cuda.empty_cache()
return max(metrics["test_accuracy"])
def optuna_optimize(n_trials=50, study_name="cifar10dvs_conv_dyned_snn_v8", n_jobs=1):
import optuna
storage = f"sqlite:///{os.path.join(OUTPUT_DIR, study_name + '.db')}"
study = optuna.create_study(
study_name=study_name,
storage=storage,
direction="maximize",
load_if_exists=True,
pruner=optuna.pruners.MedianPruner(n_startup_trials=5, n_warmup_steps=10),
)
print(f"Optuna: {n_trials} trials, {n_jobs} jobs")
print(f"DB: {storage}")
study.optimize(optuna_objective, n_trials=n_trials, n_jobs=n_jobs, gc_after_trial=True)
print(f"\nBest test accuracy: {study.best_trial.value:.4f}")
for key, value in study.best_trial.params.items():
print(f" {key}: {value}")
import json
with open(os.path.join(OUTPUT_DIR, f"{study_name}_best_params.json"), "w") as f:
json.dump({"accuracy": study.best_trial.value, "optuna_epochs": 50, **study.best_trial.params}, f, indent=2)
return study
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="CIFAR10-DVS VGGSNN+DyNED+cAdLIF+TET v8")
parser.add_argument("--optuna", action="store_true", help="Run Optuna HPO")
parser.add_argument("--n-trials", type=int, default=50, help="Number of Optuna trials")
parser.add_argument("--n-jobs", type=int, default=1, help="Parallel Optuna jobs")
parser.add_argument("--study-name", type=str, default="cifar10dvs_conv_dyned_snn_v8", help="Optuna study name")
parser.add_argument("--json", type=str, default=None, help="Best params JSON path")
parser.add_argument("--gen-data", action="store_true", help="Pre-generate cache, then exit")
parser.add_argument("--cpu", action="store_true", help="Force CPU")
parser.add_argument(
"--subset-size",
type=int,
default=None,
help="Limit train+test split to N samples (stratified by class) for quick smoke tests",
)
parser.add_argument(
"--quick-epochs", type=int, default=None, help="Override num_epochs (use with --subset-size for fast iteration)"
)
parser.add_argument(
"--dynedc-in-forward",
action="store_true",
help="Run lossless DyNEDc compress+decompress every forward pass "
"(default: off; turn on for eval-only / verifiable in-path runs)",
)
args, _ = parser.parse_known_args()
_CLI_SUBSET_SIZE = args.subset_size
_CLI_QUICK_EPOCHS = args.quick_epochs
_CLI_DYNEDC_IN_FORWARD = args.dynedc_in_forward
if args.cpu:
os.environ["CUDA_VISIBLE_DEVICES"] = ""
if args.gen_data:
assets_dir = os.path.join("..", "assets")
print("Pre-generating CIFAR10-DVS temporal frame cache...")
CIFAR10DVSFrameDataset(save_to=assets_dir, augment=False, cache=True)
print("Done.")
elif args.optuna:
optuna_optimize(n_trials=args.n_trials, study_name=args.study_name, n_jobs=args.n_jobs)
else:
main(json_path=args.json)