cifar10dvs_conv_dyned_snn_v9.py
CIFAR10-DVS DyNED SNN - 1,496 lines.
View on GitHub (image-neuro/cifar10dvs_conv_dyned_snn_v9.py).
Source
Section titled “Source”"""
CIFAR10-DVS SNN with VGGSNN Conv + DyNED Encoding + TET Loss (v9, DyNED-only)
Pipeline: CIFAR10-DVS events -> dense frames [16, 2, 32, 32] -> VGG conv per frame
-> proj_down -> DyNED (sigma-delta across time) -> proj_up -> cAdLIF SNN
v9 changes vs v8 (training-dynamics only, architecture unchanged):
- cAdLIF cell parameters (alpha_raw, beta_raw, a_raw, b_raw) and the readout
alpha_out_raw are placed in their own optimiser group with weight_decay=0.
Without this, weight decay drives `a` and `b` toward zero and degenerates
cAdLIF into a vanilla LIF, costing the network its main adaptation primitive.
- TET loss is enabled from epoch 50 (rather than epoch 200), so per-timestep
gradient pressure shapes the model from early on instead of only correcting
it at the end.
- Delays are initialised uniformly in [0, max_delay] (Hammouamri 2024) instead
of all-zeros; with zero-init every delay sat at the lower clamp boundary and
never drifted.
Run: uv run python image-neuro/cifar10dvs_conv_dyned_snn_v9.py
Optuna: uv run python image-neuro/cifar10dvs_conv_dyned_snn_v9.py --optuna --n-trials 50
Quick smoke: uv run python image-neuro/cifar10dvs_conv_dyned_snn_v9.py --subset-size 5000 --quick-epochs 30
"""
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
# 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_output_v9")
os.makedirs(OUTPUT_DIR, exist_ok=True)
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
# New caches save to v9's own folder; fall back to v8's existing cache
# (raw event-frame preprocessing is identical) so we don't have to
# rebuild 1.3 GB on first run.
cache_path = Path(save_to) / "cifar10dvs_conv_dyned_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
# =============================================================================
# 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,
):
super().__init__()
self.hidden_size = hidden_size
self.num_outputs = num_outputs
self.max_delay = max_delay
self.dyned_dim = dyned_dim
# 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
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)
# 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.
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).permute(0, 2, 1) # [B, T, dyned_dim]
encoded_flat = encoded_3d.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": []}
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)
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)
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)
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])
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 subset for fast 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)
indices = []
for cls_idx in range(n_classes):
cls_indices = np.where(labels_np == cls_idx)[0]
indices.extend(cls_indices[:per_class].tolist())
indices = sorted(set(indices))[:max_samples]
print(f"** Subset mode: stratified subset of {len(indices)} samples")
full_dataset = torch.utils.data.Subset(full_dataset, indices)
else:
full_dataset = torch.utils.data.Subset(full_dataset, list(range(max_samples)))
print(f"** Subset mode: first {max_samples} samples (no labels for stratification)")
n_total = len(full_dataset)
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 substantially.
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 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()
if metrics["layer_firing_rates"]:
fr_history = metrics["layer_firing_rates"]
keys = sorted(fr_history[0].keys())
plt.figure(figsize=(12, 6))
for key in keys:
rates = [fr[key] for fr in fr_history]
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()
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()
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()
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()
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()
# =============================================================================
# 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)
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
dyned_levels = 256
lambda_tet = 1e-3
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 {max_samples} samples")
print("=" * 80)
print("VGGSNN+DyNED(2ch,T=16)+cAdLIF+TET for CIFAR10-DVS - v9 (DyNED-only, training fixes)")
print("=" * 80)
print(f"Device: {device_str.upper()}")
if torch.cuda.is_available():
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(
f"Pipeline: DVS Frames [16\u00d72\u00d732\u00d732] \u2192 VGGSNN(2\u219264\u2192128\u2192256\u2192512) \u2192 DyNED(lvl={dyned_levels}) \u2192 cAdLIF SNN"
)
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,
).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}")
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"])
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)
# 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,
).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_v9", 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
_CLI_SUBSET_SIZE = None
_CLI_QUICK_EPOCHS = None
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="CIFAR10-DVS VGGSNN+DyNED+cAdLIF+TET v9 (DyNED-only)")
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_v9", 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="Stratified subset size for fast iteration (default: full dataset)",
)
parser.add_argument(
"--quick-epochs", type=int, default=None, help="Override num_epochs for fast iteration (default: 300)"
)
args, _ = parser.parse_known_args()
_CLI_SUBSET_SIZE = args.subset_size
_CLI_QUICK_EPOCHS = args.quick_epochs
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)