cifar10_conv_dyned_dynedc_snn_v2.py
CIFAR-10 DyNED + DyNEDc SNN - 1,807 lines.
View on GitHub (image/cifar10_conv_dyned_dynedc_snn_v2.py).
Source
Section titled “Source”"""
CIFAR-10 SNN with Conv + DyNED + DyNEDc (v2, Option B: cAdLIF + delays + TET)
Same architecture as cifar10_conv_dyned_snn_v2.py, plus a post-training
DyNEDc compression-stats pass on the cAdLIF spike outputs (Option B):
the trained network's binary spike trains from each cAdLIF layer are
losslessly compressed with DyNEDcCompressorV4 across the entire test set.
Pipeline:
Image -> Conv2d(1->64->128->256) -> FFT2 -> magnitude -> log1p
-> project -> sigma-delta (DyNED across time)
-> Linear + BN + Delay + cAdLIF (3 layers)
-> alpha-weighted readout -> TET loss
At inference (post-training):
-> DyNEDc compress per-sample on spk1, spk2, spk3
-> dynedc_compression_stats.json (per-layer ratios, savings, mode distribution)
Run: uv run python cifar10_conv_dyned_dynedc_snn_v2.py
Optuna: uv run python cifar10_conv_dyned_dynedc_snn_v2.py --optuna --n-trials 50 --n-jobs 1
"""
import argparse
import gc
import os
import sys
import time
import warnings
import math
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 snntorch import functional as SF
from snntorch import surrogate
from torchvision import datasets, transforms
try:
from sklearn.manifold import TSNE
HAS_TSNE = True
except ImportError:
HAS_TSNE = False
# DyNEDc compressor lives at the head-tracker root; add it to sys.path so this
# script (under image/) can import it without a package install.
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from dyned import DyNEDcCompressorV4 # noqa: E402
import json
from collections import Counter
# 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, "cifar10_conv_dyned_dynedc_output_v2")
os.makedirs(OUTPUT_DIR, exist_ok=True)
CIFAR10_CLASSES = None # Populated from dataset at runtime
# =============================================================================
# DyNED Encoder Layer - sigma-delta in PyTorch with STE
# =============================================================================
class DyNEDEncoderLayer(nn.Module):
"""DyNED sigma-delta quantization as a differentiable PyTorch layer.
Runs the sigma-delta feedback loop in pure PyTorch (stays on GPU).
Uses straight-through estimator: forward pass returns quantized values,
backward pass passes gradients through as if quantization was identity.
The FFT step (FFT2 -> magnitude -> log1p) is applied in the network forward
pass before features reach this layer, matching the full DyNED pipeline.
"""
def __init__(self, levels=256):
super().__init__()
self.levels = levels
def forward(self, x):
# x: [batch, features] - continuous conv activations
if not self.training:
return self._sigma_delta(x)
# STE: forward uses quantized values, backward uses x directly
quantized = self._sigma_delta(x)
return x + (quantized - x).detach()
@torch.autocast("cuda", enabled=False)
def _sigma_delta(self, x):
"""Sigma-delta quantization loop - vectorized across batch.
Forced to float32 - the sequential error accumulation over thousands of
steps loses precision in float16.
"""
x = x.float()
batch_size, n_features = x.shape
# Normalize each sample to [0, 1]
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
# Denormalize back
return quantized * x_range + x_min
# =============================================================================
# STDP + LIF Components
# =============================================================================
class EnhancedSTDP:
def __init__(self, learning_rate=0.01, time_window=20.0, a_plus=0.1, a_minus=0.12):
self.lr = learning_rate
self.tau = time_window
self.a_plus = a_plus
self.a_minus = a_minus
self.eligibility_trace = None
def compute_weight_update(self, pre_spike, post_spike, t_current):
device = pre_spike.device
t = torch.tensor(t_current, dtype=torch.float32, device=device)
if self.eligibility_trace is None or self.eligibility_trace.shape != (post_spike.size(1), pre_spike.size(1)):
self.eligibility_trace = torch.zeros(post_spike.size(1), pre_spike.size(1), device=device)
n = pre_spike.size(0)
pos_corr = (post_spike.t() @ pre_spike) / n
neg_corr = ((1 - post_spike).t() @ (1 - pre_spike)) / n
self.eligibility_trace = 0.9 * self.eligibility_trace + pos_corr
time_factor = torch.exp(-t / self.tau)
effective_lr = self.lr * time_factor
dw = (self.a_plus * pos_corr - self.a_minus * neg_corr) + 0.1 * self.eligibility_trace
dw = (dw - dw.mean()) / (dw.std() + 1e-8)
return effective_lr * torch.tanh(dw)
class EnhancedSTDPLayer(nn.Module):
def __init__(self, in_features, out_features):
super().__init__()
self.in_features = in_features
self.out_features = out_features
bound = 1 / math.sqrt(in_features)
self.weights = nn.Parameter(torch.zeros(out_features, in_features).uniform_(-bound, bound))
self.bias = nn.Parameter(torch.zeros(out_features))
self.stdp = EnhancedSTDP()
self.current_time = 0
self.activity_scale = nn.Parameter(torch.ones(1))
self.dropout = nn.Dropout(0.1)
self.pre_activations = None
self.post_activations = None
def forward(self, x):
if x.dim() == 3:
x = x.view(-1, self.in_features)
x = self.dropout(x)
output = F.linear(x, self.weights, self.bias) * torch.sigmoid(self.activity_scale)
if self.training:
self.pre_activations = x.detach()
self.post_activations = output.detach()
return output
def apply_stdp_updates(self):
if self.training and self.pre_activations is not None:
if self.pre_activations.dim() > 2:
self.pre_activations = self.pre_activations.view(-1, self.in_features)
if self.post_activations.dim() > 2:
self.post_activations = self.post_activations.view(-1, self.out_features)
dw = self.stdp.compute_weight_update(self.pre_activations, self.post_activations, self.current_time)
self.last_update_magnitude = dw.abs().mean().item()
new_weights = self.weights.data + dw
self.weights.data = torch.clamp(new_weights, -2.0, 2.0)
if hasattr(self, "bias_momentum"):
self.bias_momentum = 0.9 * self.bias_momentum + 0.1 * self.post_activations.mean(0)
else:
self.bias_momentum = self.post_activations.mean(0)
self.bias.data += 0.01 * self.bias_momentum
self.current_time += 1
self.pre_activations = None
self.post_activations = None
def reset_state(self):
self.current_time = 0
self.pre_activations = None
self.post_activations = None
if hasattr(self, "bias_momentum"):
self.bias_momentum = None
if hasattr(self.stdp, "eligibility_trace"):
self.stdp.eligibility_trace = None
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 constants (Deckers et al. 2024)
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
def apply_delays(h_seq, delays, max_delay):
"""Per-neuron fractional delays via linear interpolation between integer offsets.
h_seq: [T, B, N], delays: [N], max_delay: int
"""
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
def tet_loss(m_out, targets, label_smoothing=0.1, lambda_tet=1e-3):
"""TET loss: per-timestep CE averaged + temporal consistency MSE."""
T = m_out.shape[0]
ce_loss = sum(F.cross_entropy(m_out[t], targets, label_smoothing=label_smoothing) for t in range(T)) / T
m_mean = m_out.mean(dim=0, keepdim=True)
mse_reg = F.mse_loss(m_out, m_mean.expand_as(m_out))
return ce_loss + lambda_tet * mse_reg
# =============================================================================
# DyNEDc compression on cAdLIF spike outputs (Option B, post-training measurement)
# =============================================================================
def dynedc_compress_binary(binary_2d, chunk_size=4):
"""Compress a binary array (uint8, values in {0, 1}) with DyNEDcCompressorV4."""
flat = np.asarray(binary_2d).astype(np.uint8).flatten()
compressor = DyNEDcCompressorV4(chunk_size=chunk_size)
compressed, info = compressor.compress(flat)
return compressed, info
def measure_compression_stats(net, test_loader, device, chunk_size=4):
"""DyNEDc compression statistics on cAdLIF spike outputs across the full test set.
For every test sample, runs the network through `diagnostic_forward`, then
compresses each cAdLIF layer's [T, hidden] binary spike train independently
and records the per-sample compression ratio + mode chosen by V4.
"""
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 = 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
# =============================================================================
# Conv + DyNED + SNN Network
# =============================================================================
class ConvDyNEDSNNetwork(nn.Module):
"""Conv feature extraction -> DyNED encoding -> cAdLIF SNN classification (v2).
Pipeline:
Image [B, 1, 32, 32]
-> Conv2d (1->64->128->256) -> [B, 256, 8, 8]
-> Direct path: conv flat -> hidden_size
-> DyNED path: FFT2 -> magnitude -> log1p -> proj_down -> sigma-delta -> proj_up
-> Gated residual: encoded = (1-gate)*direct + gate*dyned_out
Time loop (T = num_steps, constant input replicated over time):
-> Linear + BN + per-neuron synaptic delay + cAdLIF (3 layers)
-> Linear readout with leaky integrator -> per-timestep membrane m_out
Returns: m_out [T, B, num_outputs] for TET loss
"""
def __init__(
self,
hidden_size=2048,
num_outputs=10,
num_steps=25,
dyned_levels=256,
dyned_dim=512,
dropout=0.2,
conv_dropout=0.1,
max_delay=5,
):
super().__init__()
self.num_steps = num_steps
self.hidden_size = hidden_size
self.num_outputs = num_outputs
self.max_delay = max_delay
# Conv front-end (unchanged from v1)
self.conv = nn.Sequential(
nn.Conv2d(1, 64, 3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.Conv2d(64, 128, 3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.MaxPool2d(2), # 16x16
nn.Conv2d(128, 256, 3, padding=1),
nn.BatchNorm2d(256),
nn.ReLU(),
nn.MaxPool2d(2), # 8x8
nn.Dropout2d(conv_dropout),
)
conv_out_size = 256 * 8 * 8 # 16384
# Direct path
self.direct_proj = nn.Linear(conv_out_size, hidden_size)
self.direct_norm = nn.LayerNorm(hidden_size)
# DyNED path
self.dyned_dim = dyned_dim
self.proj_down = nn.Linear(conv_out_size, self.dyned_dim)
self.proj_down_norm = nn.LayerNorm(self.dyned_dim)
self.dyned_encoder = DyNEDEncoderLayer(levels=dyned_levels)
self.proj_up = nn.Linear(self.dyned_dim, hidden_size)
self.proj_up_norm = nn.LayerNorm(hidden_size)
self.dyned_gate = nn.Parameter(torch.tensor([-3.0]))
# cAdLIF stack (3 hidden layers, matches CIFAR10-DVS v9 pattern)
hidden2 = hidden_size // 2
hidden3 = hidden2 // 2
self.hidden2 = hidden2
self.hidden3 = hidden3
# Layer 1
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)
# Layer 2
self.fc2 = nn.Linear(hidden_size, hidden2)
self.bn2 = nn.BatchNorm1d(hidden2, momentum=0.05)
self.cadlif2 = cAdLIFNeuron(hidden2)
self.delay2 = nn.Parameter(torch.empty(hidden2).uniform_(0.0, float(max_delay)))
self.drop2 = nn.Dropout(dropout)
# Layer 3
self.fc3 = nn.Linear(hidden2, hidden3)
self.bn3 = nn.BatchNorm1d(hidden3, momentum=0.05)
self.cadlif3 = cAdLIFNeuron(hidden3)
self.delay3 = nn.Parameter(torch.empty(hidden3).uniform_(0.0, float(max_delay)))
self.drop3 = nn.Dropout(dropout)
# Readout: linear + leaky integrator (alpha-weighted)
self.fc_out = nn.Linear(hidden3, num_outputs)
self.alpha_out_raw = nn.Parameter(torch.tensor(2.0)) # sigmoid(2.0) ~ 0.88
# Kaiming init for FC weights
for fc in (self.fc1, self.fc2, self.fc3):
nn.init.kaiming_uniform_(fc.weight)
def _get_alpha_out(self):
return torch.sigmoid(self.alpha_out_raw)
def _encode(self, x):
"""Compute the gated-residual encoded features [B, hidden_size]."""
conv_out = self.conv(x)
conv_flat = conv_out.flatten(1)
direct = self.direct_norm(self.direct_proj(conv_flat))
with torch.amp.autocast("cuda", enabled=False):
fft_result = torch.fft.fft2(conv_out.float())
fft_shifted = torch.fft.fftshift(fft_result, dim=(-2, -1))
fft_magnitude = torch.log1p(torch.abs(fft_shifted))
compact = self.proj_down_norm(self.proj_down(fft_magnitude.flatten(1)))
encoded_compact = self.dyned_encoder(compact)
dyned_out = self.proj_up_norm(self.proj_up(encoded_compact))
gate = torch.sigmoid(self.dyned_gate)
return (1 - gate) * direct + gate * dyned_out
def forward(self, x, return_temporal=False):
# x: [B, 1, 32, 32] grayscale CIFAR-10
B = x.size(0)
T = self.num_steps
device = x.device
encoded = self._encode(x) # [B, hidden]
encoded_seq = encoded.unsqueeze(0).expand(T, B, self.hidden_size)
# ---- Layer 1: precompute Linear + BN + delays, then unroll cAdLIF ----
h1 = self.fc1(encoded_seq) # [T, B, hidden]
h1 = self.bn1(h1.reshape(-1, self.hidden_size)).reshape(T, B, self.hidden_size)
h1 = apply_delays(h1, self.delay1.clamp(0.0, float(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 = torch.stack(s1_list)
s1_seq = self.drop1(s1_seq)
# ---- Layer 2 ----
h2 = self.fc2(s1_seq)
h2 = self.bn2(h2.reshape(-1, self.hidden2)).reshape(T, B, self.hidden2)
h2 = apply_delays(h2, self.delay2.clamp(0.0, float(self.max_delay)), self.max_delay)
s2_list = []
u2 = torch.zeros(B, self.hidden2, device=device)
w2 = torch.zeros(B, self.hidden2, device=device)
s2 = torch.zeros(B, self.hidden2, device=device)
for t in range(T):
s2, u2, w2 = self.cadlif2(h2[t], u2, w2, s2)
s2_list.append(s2)
s2_seq = torch.stack(s2_list)
s2_seq = self.drop2(s2_seq)
# ---- Layer 3 ----
h3 = self.fc3(s2_seq)
h3 = self.bn3(h3.reshape(-1, self.hidden3)).reshape(T, B, self.hidden3)
h3 = apply_delays(h3, self.delay3.clamp(0.0, float(self.max_delay)), self.max_delay)
s3_list = []
u3 = torch.zeros(B, self.hidden3, device=device)
w3 = torch.zeros(B, self.hidden3, device=device)
s3 = torch.zeros(B, self.hidden3, device=device)
for t in range(T):
s3, u3, w3 = self.cadlif3(h3[t], u3, w3, s3)
s3_list.append(s3)
s3_seq = torch.stack(s3_list)
s3_seq = self.drop3(s3_seq)
# ---- Readout: alpha-weighted membrane potential ----
alpha_out = self._get_alpha_out()
u_out = torch.zeros(B, self.num_outputs, device=device)
m_out = []
for t in range(T):
cur = self.fc_out(s3_seq[t])
u_out = alpha_out * u_out + cur
m_out.append(u_out)
m_out = torch.stack(m_out) # [T, B, num_outputs]
if return_temporal:
return m_out, {
"spk1": s1_seq.detach(),
"spk2": s2_seq.detach(),
"spk3": s3_seq.detach(),
"m_out": m_out.detach(),
}
return m_out
def diagnostic_forward(self, x):
"""Same as forward(return_temporal=True), but also returns encoded features."""
encoded = self._encode(x)
m_out, layer_data = self.forward(x, return_temporal=True)
layer_data["encoded"] = encoded.detach()
layer_data["m_out"] = m_out.detach()
return layer_data
def get_param_groups(self, lr, lr_delay, weight_decay):
"""Optimizer param groups: cAdLIF cell params + alpha_out_raw with weight_decay=0,
delays on a separate (typically smaller) learning rate, everything else
on the main lr+wd group."""
cadlif_params, delay_params, no_decay_other, regular_params = [], [], [], []
for name, p in self.named_parameters():
if not p.requires_grad:
continue
if any(k in name for k in (".alpha_raw", ".beta_raw", ".a_raw", ".b_raw")):
cadlif_params.append(p)
elif "delay" in name and "delay" == name.split(".")[-1].rstrip("0123456789"):
# match delay1/delay2/delay3 only
delay_params.append(p)
elif name.endswith("alpha_out_raw"):
no_decay_other.append(p)
else:
regular_params.append(p)
return [
{"params": regular_params, "lr": lr, "weight_decay": weight_decay},
{"params": cadlif_params, "lr": lr, "weight_decay": 0.0},
{"params": no_decay_other, "lr": lr, "weight_decay": 0.0},
{"params": delay_params, "lr": lr_delay, "weight_decay": 0.0},
]
# =============================================================================
# Loss Function
# =============================================================================
def enhanced_combined_loss(spk_rec, mem_rec, targets, current_epoch):
rate_loss = SF.ce_rate_loss()(spk_rec, targets)
mem_avg = mem_rec.mean(0)
mem_loss = F.cross_entropy(mem_avg, targets)
spk_diff = spk_rec[1:] - spk_rec[:-1]
temporal_loss = torch.mean(torch.abs(spk_diff))
firing_rates = torch.mean(spk_rec, dim=(0, 1))
target_rates = torch.full_like(firing_rates, 0.2)
sparsity_loss = F.mse_loss(firing_rates, target_rates)
spike_rates = spk_rec.mean(0)
centroids = []
for cls in range(spike_rates.size(1)):
mask = targets == cls
if mask.sum() > 1:
centroids.append(spike_rates[mask].mean(0))
if len(centroids) >= 2:
centroids = torch.stack(centroids)
centroids_norm = F.normalize(centroids, dim=1)
similarity = torch.mm(centroids_norm, centroids_norm.t())
eye_mask = ~torch.eye(len(centroids), dtype=torch.bool, device=similarity.device)
cosine_loss = similarity[eye_mask].mean()
else:
cosine_loss = torch.tensor(0.0, device=spk_rec.device)
progress = min(current_epoch / 100, 1.0)
rate_weight = 0.6 + 0.2 * progress
mem_weight = 0.2 * (1 - progress)
temporal_weight = 0.1 * (1 - progress)
sparsity_weight = 0.1
cosine_weight = 0.1
total_loss = (
rate_weight * rate_loss
+ mem_weight * mem_loss
+ temporal_weight * temporal_loss
+ sparsity_weight * sparsity_loss
+ cosine_weight * cosine_loss
)
return total_loss
# =============================================================================
# Data Setup - standard transforms (DyNED is in the network, not the transform)
# =============================================================================
def setup_training(batch_size=128, workers=4, pin_memory=True, multiprocessing_context=None):
global CIFAR10_CLASSES
train_transform = transforms.Compose(
[
transforms.RandomAffine(degrees=15, translate=(0.1, 0.1), scale=(0.9, 1.1)),
transforms.RandomPerspective(distortion_scale=0.2, p=0.5),
transforms.ToTensor(),
transforms.Grayscale(num_output_channels=1),
transforms.RandomErasing(p=0.1),
transforms.Normalize((0.5,), (0.5,)),
]
)
test_transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Grayscale(num_output_channels=1),
transforms.Normalize((0.5,), (0.5,)),
]
)
train_dataset = datasets.CIFAR10(root="../assets", train=True, download=True, transform=train_transform)
CIFAR10_CLASSES = train_dataset.classes
test_dataset = datasets.CIFAR10(root="../assets", train=False, download=True, transform=test_transform)
loader_kwargs = dict(pin_memory=pin_memory)
if workers > 0:
loader_kwargs["persistent_workers"] = False
loader_kwargs["prefetch_factor"] = 4
if multiprocessing_context is not None:
loader_kwargs["multiprocessing_context"] = multiprocessing_context
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=workers,
**loader_kwargs,
)
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=batch_size * 2,
shuffle=False,
num_workers=workers,
**loader_kwargs,
)
return train_loader, test_loader
# =============================================================================
# Training Loop
# =============================================================================
def train_network(
net,
train_loader,
test_loader,
num_epochs=250,
device="cuda",
lr=2e-3,
lr_delay=1e-2,
weight_decay=0.01,
pct_start=0.3,
lambda_tet=1e-3,
tet_start_epoch=50,
trial=None,
):
"""v2 training loop: TET loss schedule + optimizer groups + cAdLIF-friendly clipping.
For epochs < tet_start_epoch, uses mean cross-entropy on the time-averaged
readout (faster initial convergence). From tet_start_epoch onward, switches
to per-timestep TET loss for stronger temporal credit assignment.
"""
param_groups = net.get_param_groups(lr=lr, lr_delay=lr_delay, weight_decay=weight_decay)
optimizer = torch.optim.AdamW(param_groups, betas=(0.9, 0.999))
scheduler = torch.optim.lr_scheduler.OneCycleLR(
optimizer,
max_lr=[g["lr"] for g in param_groups],
epochs=num_epochs,
steps_per_epoch=len(train_loader),
pct_start=pct_start,
anneal_strategy="cos",
div_factor=25,
final_div_factor=1000,
)
scaler = torch.amp.GradScaler("cuda")
best_acc = 0
metrics = {
"epoch": [],
"train_loss": [],
"test_accuracy": [],
"learning_rate": [],
"epoch_time": [],
"layer_firing_rates": [],
"per_class_accuracy": [],
"batch_losses": [],
}
diag_batch = next(iter(test_loader))
batch_size = train_loader.batch_size
print(f"Training on device: {device}")
print(f"Batch size: {batch_size}, num_epochs: {num_epochs}, tet_start_epoch: {tet_start_epoch}")
if device == "cuda":
print(
f"GPU memory before training: {torch.cuda.memory_allocated() / 1024**3:.2f} GB allocated, "
f"{torch.cuda.memory_reserved() / 1024**3:.2f} GB reserved"
)
for epoch in range(num_epochs):
epoch_start = time.time()
net.train()
running_loss = 0.0
epoch_loss = 0.0
num_batches = 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()
with torch.amp.autocast("cuda"):
m_out = net(data) # [T, B, num_outputs]
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)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
scheduler.step()
batch_loss = loss.item()
running_loss += batch_loss
epoch_loss += batch_loss
num_batches += 1
metrics["batch_losses"].append(batch_loss)
if i % 100 == 99:
avg_loss = running_loss / 100
print(f"Epoch {epoch + 1}, Batch {i + 1}: Loss = {avg_loss:.4f}", end="")
if device == "cuda":
print(f" | GPU Memory: {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"]
with torch.no_grad():
diag_data, _ = diag_batch
diag_data = diag_data[:32].to(device)
net.eval()
layer_data = net.diagnostic_forward(diag_data)
net.train()
firing_rates = {
key.replace("spk", "layer"): layer_data[key].mean().item()
for key in layer_data
if key.startswith("spk")
}
eval_result = evaluate(net, test_loader, device)
test_acc = eval_result["accuracy"]
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"])
loss_label = "TET" if use_tet else "CE"
print(
f"Epoch {epoch + 1}: Test Accuracy = {test_acc:.4f} | "
f"{loss_label} Loss = {avg_epoch_loss:.4f} | LR = {current_lr:.6f} | "
f"FR = [{firing_rates.get('layer1', 0):.3f}, {firing_rates.get('layer2', 0):.3f}, {firing_rates.get('layer3', 0):.3f}] | "
f"Time = {epoch_time:.1f}s"
)
if device == "cuda":
torch.cuda.empty_cache()
if test_acc > best_acc:
best_acc = test_acc
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_conv_dyned_model.pth"),
)
# Optuna pruning support
if trial is not None:
import optuna
trial.report(test_acc, epoch)
if trial.should_prune():
raise optuna.TrialPruned()
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)
m_out = net(data) # [T, B, num_outputs]
pred = m_out.mean(0) # time-averaged readout
predicted_classes = pred.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(pred.cpu())
accuracy = correct / total
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": {CIFAR10_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 visualize_dyned_effect(net, sample_batch, device):
"""Visualize how DyNED transforms conv features (FFT + sigma-delta)."""
net.eval()
with torch.no_grad():
data, targets = sample_batch
data = data[:8].to(device)
# Get conv features -> FFT -> project -> sigma-delta
conv_out = net.conv(data)
with torch.amp.autocast("cuda", enabled=False):
fft_result = torch.fft.fft2(conv_out.float())
fft_shifted = torch.fft.fftshift(fft_result, dim=(-2, -1))
fft_magnitude = torch.log1p(torch.abs(fft_shifted))
compact = net.proj_down_norm(net.proj_down(fft_magnitude.flatten(1)))
encoded = net.dyned_encoder(compact)
conv_features = compact # show the compact FFT magnitude representation
conv_np = conv_features.cpu().numpy()
enc_np = encoded.cpu().numpy()
fig, axes = plt.subplots(3, 2, figsize=(16, 12))
# Feature distributions
axes[0, 0].hist(conv_np.flatten(), bins=100, alpha=0.7, color="steelblue", edgecolor="none")
axes[0, 0].set_title("Conv Features Distribution")
axes[0, 0].set_xlabel("Activation value")
axes[0, 1].hist(enc_np.flatten(), bins=100, alpha=0.7, color="coral", edgecolor="none")
axes[0, 1].set_title("DyNED-Encoded Distribution")
axes[0, 1].set_xlabel("Encoded value")
# Per-sample comparison (first 4 samples)
for i in range(min(4, len(data))):
ax = axes[1, 0] if i < 2 else axes[2, 0]
ax_enc = axes[1, 1] if i < 2 else axes[2, 1]
alpha = 0.7 if i % 2 == 0 else 0.4
ax.plot(conv_np[i, :200], alpha=alpha, linewidth=0.5, label=f"Sample {i} ({CIFAR10_CLASSES[targets[i]]})")
ax_enc.plot(
enc_np[i, :200], alpha=alpha, linewidth=0.5, label=f"Sample {i} ({CIFAR10_CLASSES[targets[i]]})"
)
for ax in [axes[1, 0], axes[2, 0]]:
ax.set_title("Conv Features (first 200 dims)")
ax.legend(fontsize=7)
for ax in [axes[1, 1], axes[2, 1]]:
ax.set_title("DyNED Encoded (first 200 dims)")
ax.legend(fontsize=7)
# Quantization error
error = conv_np - enc_np
fig.suptitle(
f"DyNED Encoding Effect - levels={net.dyned_encoder.levels}, mean error={np.abs(error).mean():.4f}",
fontsize=14,
)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "conv_dyned_encoding_effect.png"), dpi=150)
plt.close()
print("Saved DyNED encoding effect visualization")
def visualize_network_activity(net, sample_batch, device):
net.eval()
with torch.no_grad():
data, targets = sample_batch
data = data.to(device)
spk_rec, mem_rec = net(data)
fig, axes = plt.subplots(2, 2, figsize=(15, 15))
axes[0, 0].imshow(spk_rec[:, 0].cpu().numpy(), aspect="auto", cmap="binary")
axes[0, 0].set_title("Spike Raster (First Sample)")
axes[0, 0].set_xlabel("Neuron Index")
axes[0, 0].set_ylabel("Time Step")
axes[0, 1].plot(mem_rec[:, 0].cpu().numpy())
axes[0, 1].set_title("Membrane Potential (First Sample)")
axes[0, 1].set_xlabel("Time Step")
axes[0, 1].set_ylabel("Membrane Potential")
mean_rate = spk_rec.mean(0).cpu().numpy()
axes[1, 0].hist(mean_rate.flatten(), bins=50)
axes[1, 0].set_title("Average Firing Rate Distribution")
axes[1, 0].set_xlabel("Firing Rate")
axes[1, 0].set_ylabel("Count")
# Visualize learned conv filters (first layer)
filters = net.conv[0].weight.data.cpu()
n_show = min(16, filters.size(0))
grid = filters[:n_show, 0].numpy()
axes[1, 1].imshow(
np.concatenate([grid[i] for i in range(n_show)], axis=1),
cmap="gray",
aspect="auto",
)
axes[1, 1].set_title(f"Learned Conv1 Filters (first {n_show})")
axes[1, 1].set_xlabel("Filter (concatenated)")
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "conv_dyned_network_activity.png"), dpi=150)
plt.close()
def visualize_conv_features(net, sample_batch, device):
"""Visualize conv layer activations for a sample image."""
net.eval()
with torch.no_grad():
data, targets = sample_batch
img = data[:1].to(device)
activations = []
x = img
for layer in net.conv:
x = layer(x)
if isinstance(layer, nn.Conv2d):
activations.append(x.cpu())
n_layers = len(activations)
fig, axes = plt.subplots(1, n_layers + 1, figsize=(5 * (n_layers + 1), 5))
axes[0].imshow(img[0, 0].cpu().numpy(), cmap="gray")
axes[0].set_title("Original")
for i, act in enumerate(activations):
axes[i + 1].imshow(act[0].mean(0).numpy(), cmap="viridis")
axes[i + 1].set_title(f"Conv {i + 1} ({act.shape[1]} ch, {act.shape[2]}x{act.shape[3]})")
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "conv_dyned_conv_features.png"), dpi=150)
plt.close()
def save_unified_metrics(metrics):
filepath = os.path.join(OUTPUT_DIR, "conv_dyned_training_metrics.csv")
header_parts = ["epoch", "train_loss", "test_accuracy", "learning_rate", "epoch_time"]
fr_keys = []
if metrics["layer_firing_rates"]:
fr_keys = sorted(metrics["layer_firing_rates"][0].keys())
header_parts.extend(f"firing_rate_{k}" for k in fr_keys)
header_parts.extend(f"acc_{cls}" for cls in CIFAR10_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 CIFAR10_CLASSES:
row.append(pca.get(cls, 0))
rows.append(row)
np.savetxt(filepath, np.array(rows), delimiter=",", header=header, comments="")
print(f"Saved unified metrics to {filepath}")
def visualize_layer_spike_rasters(layer_data):
fig, axes = plt.subplots(2, 2, figsize=(16, 12))
layer_names = [
("spk1", "Layer 1 (Hidden)"),
("spk2", "Layer 2 (Hidden)"),
("spk3", "Layer 3 (Hidden)"),
("spk_out", "Output Layer"),
]
for ax, (key, title) in zip(axes.flat, layer_names):
spikes = layer_data[key][:, 0].cpu().numpy()
ax.imshow(spikes, aspect="auto", cmap="binary", interpolation="nearest")
ax.set_title(title)
ax.set_xlabel("Neuron Index")
ax.set_ylabel("Time Step")
rate = spikes.mean()
ax.text(
0.02,
0.98,
f"Rate: {rate:.3f}",
transform=ax.transAxes,
va="top",
fontsize=9,
color="red",
bbox=dict(boxstyle="round", facecolor="white", alpha=0.8),
)
plt.suptitle("Per-Layer Spike Rasters (Sample 0)", fontsize=14)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "conv_dyned_layer_spike_rasters.png"), dpi=150)
plt.close()
print("Saved layer spike rasters")
def visualize_membrane_distributions(layer_data):
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
layer_names = [
("mem1", "Layer 1"),
("mem2", "Layer 2"),
("mem3", "Layer 3"),
("mem_out", "Output Layer"),
]
for ax, (key, title) in zip(axes.flat, layer_names):
mem = layer_data[key].cpu().numpy().flatten()
ax.hist(mem, bins=100, density=True, alpha=0.7, color="steelblue")
ax.set_title(f"{title} Membrane Potential Distribution")
ax.set_xlabel("Membrane Potential")
ax.set_ylabel("Density")
ax.axvline(x=0.5, color="red", linestyle="--", label="Threshold")
ax.legend()
ax.text(
0.02,
0.98,
f"Mean: {mem.mean():.3f}\nStd: {mem.std():.3f}",
transform=ax.transAxes,
va="top",
fontsize=9,
bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.8),
)
plt.suptitle("Membrane Potential Distributions", fontsize=14)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "conv_dyned_membrane_distributions.png"), dpi=150)
plt.close()
print("Saved membrane distributions")
def visualize_firing_rates_history(metrics):
if not metrics["layer_firing_rates"]:
return
fr_history = metrics["layer_firing_rates"]
keys = sorted(fr_history[0].keys())
epochs = metrics["epoch"]
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 Over Training")
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()
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()
def visualize_weight_distributions(initial_weights, final_net):
final_weights = {}
for name, layer in [("stdp1", final_net.stdp1), ("stdp2", final_net.stdp2), ("stdp3", final_net.stdp3)]:
final_weights[name] = layer.weights.data.cpu().numpy().flatten()
fig, axes = plt.subplots(1, 3, figsize=(18, 5))
for ax, name in zip(axes, ["stdp1", "stdp2", "stdp3"]):
ax.hist(initial_weights[name], bins=100, alpha=0.5, label="Before", density=True, color="blue")
ax.hist(final_weights[name], bins=100, alpha=0.5, label="After", density=True, color="red")
ax.set_title(f"{name} Weight Distribution")
ax.set_xlabel("Weight Value")
ax.set_ylabel("Density")
ax.legend()
plt.suptitle("Weight Distributions: Before vs After Training", fontsize=14)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "conv_dyned_weight_distributions.png"), dpi=150)
plt.close()
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=CIFAR10_CLASSES,
yticklabels=CIFAR10_CLASSES,
ylabel="True Label",
xlabel="Predicted Label",
title="Confusion Matrix (Normalized)",
)
plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
thresh = cm_normalized.max() / 2.0
for i in range(10):
for j in range(10):
ax.text(
j,
i,
f"{cm_normalized[i, j]:.2f}",
ha="center",
va="center",
color="white" if cm_normalized[i, j] > thresh else "black",
fontsize=7,
)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "confusion_matrix.png"), dpi=150)
plt.close()
np.savetxt(
os.path.join(OUTPUT_DIR, "conv_dyned_confusion_matrix.csv"),
cm,
delimiter=",",
fmt="%d",
header=",".join(CIFAR10_CLASSES),
comments="",
)
print("Saved confusion matrix")
def visualize_tsne(representations, targets):
if not HAS_TSNE:
print("Skipping t-SNE: sklearn not installed")
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 (this may take a moment)...")
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(CIFAR10_CLASSES)
plt.title("t-SNE of Output Layer Representations")
plt.xlabel("t-SNE 1")
plt.ylabel("t-SNE 2")
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "tsne.png"), dpi=150)
plt.close()
print("Saved t-SNE visualization")
def visualize_per_class_spikes(layer_data, targets):
spk_out = layer_data["spk_out"]
if isinstance(targets, torch.Tensor):
targets = targets.cpu()
fig, axes = plt.subplots(2, 5, figsize=(20, 8))
for cls_idx, ax in enumerate(axes.flat):
mask = targets == cls_idx
if mask.sum() == 0:
ax.set_title(f"{CIFAR10_CLASSES[cls_idx]} (no samples)")
continue
class_spikes = spk_out[:, mask, :].mean(dim=1).cpu().numpy()
ax.imshow(class_spikes, aspect="auto", cmap="hot", interpolation="nearest")
ax.set_title(CIFAR10_CLASSES[cls_idx])
ax.set_xlabel("Neuron")
ax.set_ylabel("Time Step")
plt.suptitle("Per-Class Average Spike Patterns (Output Layer)", fontsize=14)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "conv_dyned_per_class_spikes.png"), dpi=150)
plt.close()
print("Saved per-class spike patterns")
def visualize_eligibility_traces(net):
layers = [
("stdp1", net.stdp1),
("stdp2", net.stdp2),
("stdp3", net.stdp3),
]
traces = []
names = []
for name, layer in layers:
if layer.stdp.eligibility_trace is not None:
traces.append(layer.stdp.eligibility_trace.cpu().numpy())
names.append(name)
if not traces:
print("No eligibility traces available (need to run training first)")
return
fig, axes = plt.subplots(1, len(traces), figsize=(6 * len(traces), 5))
if len(traces) == 1:
axes = [axes]
for ax, trace, name in zip(axes, traces, names):
if trace.shape[0] > 100 or trace.shape[1] > 100:
trace = trace[:100, :100]
im = ax.imshow(trace, aspect="auto", cmap="RdBu_r", interpolation="nearest")
ax.set_title(f"{name} Eligibility Trace")
ax.set_xlabel("Pre-synaptic")
ax.set_ylabel("Post-synaptic")
plt.colorbar(im, ax=ax)
plt.suptitle("STDP Eligibility Traces", fontsize=14)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "conv_dyned_eligibility_traces.png"), dpi=150)
plt.close()
print("Saved eligibility traces")
def analyze_training_metrics(metrics):
losses = np.array(metrics["train_loss"])
accuracies = np.array(metrics["test_accuracy"])
epochs = np.array(metrics["epoch"])
save_unified_metrics(metrics)
loss_improvement = (losses[0] - losses[-1]) / losses[0] * 100
accuracy_improvement = (accuracies[-1] - accuracies[0]) / (accuracies[0] + 1e-8) * 100
best_epoch = np.argmax(accuracies)
conv_threshold = 0.01
converged_epoch = None
for i in range(1, len(losses)):
if abs((losses[i] - losses[i - 1]) / (losses[i - 1] + 1e-8)) < conv_threshold:
converged_epoch = i
break
print("\nDetailed Training Analysis:")
print(f"Total loss improvement: {loss_improvement:.2f}%")
print(f"Total accuracy improvement: {accuracy_improvement:.2f}%")
print(f"Best performing epoch: {best_epoch + 1}")
if converged_epoch:
print(f"Training converged at epoch: {converged_epoch}")
print("\nPer-class accuracy at best epoch:")
best_pca = metrics["per_class_accuracy"][best_epoch]
for cls, acc in sorted(best_pca.items(), key=lambda x: x[1]):
print(f" {cls:>12s}: {acc:.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", 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("CIFAR-10 Conv+DyNED+$\\mathrm{DyNED}_c$+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["learning_rate"]:
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()
return {
"loss_improvement": loss_improvement,
"accuracy_improvement": accuracy_improvement,
"best_epoch": best_epoch,
"converged_epoch": converged_epoch,
}
# =============================================================================
# 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.backends.cudnn.deterministic = False
torch.set_float32_matmul_precision("high")
if hasattr(torch.nn.functional, "scaled_dot_product_attention"):
torch.backends.cuda.enable_mem_efficient_sdp(True)
torch.backends.cuda.enable_flash_sdp(True)
# Hyperparameters - defaults or from Optuna JSON
num_epochs = 250
tet_start_epoch = 50
if json_path:
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", 1e-2)
weight_decay = params["weight_decay"]
hidden_size = params["hidden_size"]
dyned_dim = params.get("dyned_dim", 512)
max_delay = params.get("max_delay", 5)
num_steps = params["num_steps"]
dropout = params.get("dropout", 0.2)
conv_dropout = params.get("conv_dropout", 0.1)
batch_size = params.get("batch_size", 256)
optuna_epochs = params.get("optuna_epochs", 50)
pct_start_raw = params.get("pct_start", 0.3)
pct_start = pct_start_raw * optuna_epochs / num_epochs # scale warmup for longer training
dyned_levels = params.get("dyned_levels", 256)
lambda_tet = params.get("lambda_tet", 1e-3)
print(f" pct_start scaled: {pct_start_raw:.3f} ({optuna_epochs}ep) -> {pct_start:.3f} ({num_epochs}ep)")
else:
lr = 0.05
lr_delay = 1e-2
weight_decay = 5e-4
hidden_size = 2048
dyned_dim = 512
max_delay = 5
num_steps = 25
dropout = 0.2
conv_dropout = 0.1
batch_size = 256
pct_start = 0.3
dyned_levels = 256
lambda_tet = 1e-3
num_outputs = 10
workers = min(8, os.cpu_count() - 2) if os.cpu_count() > 2 else 0
print("=" * 80)
print("TRAINING CONFIGURATION - Conv+DyNED+cAdLIF+TET for CIFAR-10 (v2)")
print("=" * 80)
print(f"Device: {device_str.upper()}")
if torch.cuda.is_available():
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"VRAM Available: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
print(
f"Pipeline: Image -> Conv(1->64->128->256) -> proj({dyned_dim}) -> DyNED(levels={dyned_levels}) -> proj({hidden_size}) -> 3x cAdLIF + delays -> readout"
)
print(f"Hidden sizes: {hidden_size} -> {hidden_size // 2} -> {hidden_size // 4} -> {num_outputs}")
print(f"Batch size: {batch_size}")
print(f"Num steps (T): {num_steps}")
print(f"DyNED levels: {dyned_levels}")
print(f"Max delay: {max_delay}")
print(f"TET start epoch: {tet_start_epoch} of {num_epochs} (mean CE before, TET after)")
print(f"Data workers: {workers}")
print(f"Mixed precision: Enabled (FP16 + TF32)")
print("=" * 80)
print("\nInitializing components...")
warnings.filterwarnings("ignore", category=np.exceptions.VisibleDeprecationWarning, message=".*dtype.*align.*")
warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*")
net = ConvDyNEDSNNetwork(
hidden_size=hidden_size,
num_outputs=num_outputs,
num_steps=num_steps,
dyned_levels=dyned_levels,
dyned_dim=dyned_dim,
dropout=dropout,
conv_dropout=conv_dropout,
max_delay=max_delay,
).to(device)
print("Setting up data loaders...")
train_loader, test_loader = setup_training(batch_size=batch_size, workers=workers)
print("Starting training...")
try:
if device_str == "cuda":
torch.cuda.empty_cache()
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,
pct_start=pct_start,
lambda_tet=lambda_tet,
tet_start_epoch=tet_start_epoch,
)
# Post-training analysis (visualizations from v1 are STDP/LIF-specific
# and were removed in v2; only training-curve + confusion-matrix +
# per-class accuracy plots remain).
print("\n" + "=" * 80)
print("GENERATING POST-TRAINING METRICS")
print("=" * 80)
analyze_training_metrics(metrics)
visualize_firing_rates_history(metrics)
print("\nRunning final evaluation...")
final_eval = evaluate(net, test_loader, device, collect_representations=True)
visualize_confusion_matrix_plot(final_eval["confusion_matrix"])
visualize_per_class_accuracy(final_eval["per_class_accuracy"])
if "representations" in final_eval:
visualize_tsne(final_eval["representations"], final_eval["targets"])
print("\nFinal Per-Class Accuracy:")
for cls, acc in sorted(final_eval["per_class_accuracy"].items(), key=lambda x: x[1], reverse=True):
print(f" {cls:>12s}: {acc:.4f}")
torch.save(
{
"model_state_dict": net.state_dict(),
"dyned_levels": dyned_levels,
"metrics": {
"train_losses": metrics["train_loss"],
"test_accuracies": metrics["test_accuracy"],
"final_accuracy": metrics["test_accuracy"][-1],
"per_class_accuracy": metrics["per_class_accuracy"][-1],
},
},
os.path.join(OUTPUT_DIR, "conv_dyned_final_model.pth"),
)
print(f"\nTraining completed! Final accuracy: {metrics['test_accuracy'][-1]:.4f}")
print(f"Best test accuracy: {max(metrics['test_accuracy']):.4f}")
print(f"All outputs saved to: {OUTPUT_DIR}")
# ----- DyNEDc compression measurement on cAdLIF spike outputs -----
print("\n" + "=" * 80)
print("DYNEDC COMPRESSION MEASUREMENT (Option B, post-training)")
print("=" * 80)
compression_stats = measure_compression_stats(net, test_loader, device, chunk_size=4)
stats_path = os.path.join(OUTPUT_DIR, "dynedc_compression_stats.json")
with open(stats_path, "w") as f:
json.dump(compression_stats, f, indent=2)
print(f"Saved: {stats_path}")
for layer in ("spk1", "spk2", "spk3", "overall"):
s = compression_stats[layer]
print(
f" {layer:>7s}: ratio={s['mean_compression_ratio']:.4f} "
f"saving={s['mean_space_saving_pct']:.2f}% "
f"n={s['n_samples']}"
)
except KeyboardInterrupt:
print("Training interrupted! Saving current state...")
torch.save(
{
"model_state_dict": net.state_dict(),
"dyned_levels": dyned_levels,
},
os.path.join(OUTPUT_DIR, "conv_dyned_interrupted_model.pth"),
)
# =============================================================================
# Optuna Hyperparameter Optimization
# =============================================================================
def optuna_objective(trial):
"""Optuna objective for CIFAR-10 Conv+DyNED+SNN hyperparameter optimization."""
import optuna
device_str = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device_str)
# --- Suggest hyperparameters (cAdLIF + delays + TET search space) ---
lr = trial.suggest_float("lr", 1e-3, 1e-1, log=True)
lr_delay = trial.suggest_float("lr_delay", 1e-3, 1e-1, log=True)
weight_decay = trial.suggest_float("weight_decay", 1e-5, 1e-3, log=True)
hidden_size = trial.suggest_categorical("hidden_size", [1024, 2048])
dyned_dim = trial.suggest_categorical("dyned_dim", [256, 512, 1024])
max_delay = trial.suggest_categorical("max_delay", [2, 3, 5, 7])
num_steps = trial.suggest_int("num_steps", 10, 25, step=5)
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", [128, 256, 512])
pct_start = trial.suggest_float("pct_start", 0.1, 0.4)
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)
num_epochs = 50 # short runs for HPO
# 4 worker processes (spawn) move the per-batch tensor handling off the
# GIL-pegged main thread, which is otherwise busy submitting cAdLIF
# temporal-loop CUDA kernels. spawn is required because fork + an already-
# initialised CUDA context corrupts the workers.
workers = 4
net = ConvDyNEDSNNetwork(
hidden_size=hidden_size,
num_outputs=10,
num_steps=num_steps,
dyned_levels=dyned_levels,
dyned_dim=dyned_dim,
dropout=dropout,
conv_dropout=conv_dropout,
max_delay=max_delay,
).to(device)
# pin_memory=False under Optuna: pinned host memory accumulates across
# trials in the same process (empty_cache only frees CUDA-side memory),
# which eventually causes pin_memory() to fail with cudaErrorInvalidValue
# under high host-side memory pressure.
train_loader, test_loader = setup_training(
batch_size=batch_size,
workers=workers,
pin_memory=False,
multiprocessing_context="spawn" if workers > 0 else None,
)
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,
pct_start=pct_start,
lambda_tet=lambda_tet,
tet_start_epoch=10, # earlier TET start because total budget is 50 epochs
trial=trial,
)
except optuna.TrialPruned:
raise
finally:
# Tear down workers + datasets fully so the next trial starts clean.
for ld in (train_loader, test_loader):
it = getattr(ld, "_iterator", None)
if it is not None and hasattr(it, "_shutdown_workers"):
it._shutdown_workers()
del train_loader, test_loader
del net
gc.collect()
if device_str == "cuda":
torch.cuda.empty_cache()
torch.cuda.synchronize()
best_acc = max(metrics["test_accuracy"])
return best_acc
def optuna_optimize(n_trials=50, study_name="cifar10_conv_dyned_dynedc_snn_v2", n_jobs=1):
"""Run Optuna hyperparameter optimization."""
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"Starting Optuna optimization: {n_trials} trials, {n_jobs} parallel jobs")
print(f"Study database: {storage}")
study.optimize(optuna_objective, n_trials=n_trials, n_jobs=n_jobs, gc_after_trial=True)
print("\n" + "=" * 80)
print("OPTUNA OPTIMIZATION RESULTS")
print("=" * 80)
print(f"Best trial accuracy: {study.best_trial.value:.4f}")
print("Best hyperparameters:")
for key, value in study.best_trial.params.items():
print(f" {key}: {value}")
# Save best params
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="CIFAR-10 Conv+DyNED+SNN")
parser.add_argument("--optuna", action="store_true", help="Run Optuna hyperparameter optimization")
parser.add_argument("--n-trials", type=int, default=50, help="Number of Optuna trials")
parser.add_argument("--n-jobs", type=int, default=1, help="Number of parallel Optuna trials")
parser.add_argument("--study-name", type=str, default="cifar10_conv_dyned_dynedc_snn_v2", help="Optuna study name")
parser.add_argument("--json", type=str, default=None, help="Path to Optuna best params JSON file")
parser.add_argument("--cpu", action="store_true", help="Force CPU (disable CUDA)")
args, _ = parser.parse_known_args()
if args.cpu:
os.environ["CUDA_VISIBLE_DEVICES"] = ""
if 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)