speech_cadlif_dyned_snn_cached-4.py
Speech Commands DyNED SNN - 1,774 lines.
View on GitHub (speech/speech_cadlif_dyned_snn_cached-4.py).
Source
Section titled “Source”"""
Speech Commands SNN with DyNED Encoding - cAdLIF + Learnable Delays (v4)
Pipeline: Waveform -> Mel+MFCC -> per-frame DyNED (sigma-delta -> quantised features) -> cAdLIF SNN
Changes from v3 (stuck at 89-90%):
- Replace raw 513-bin STFT with 80 mel + 80 MFCC = 160 features through DyNED
Mel-scale features are perceptually weighted and more informative for speech
than raw STFT bins.
- hop_length=80 (from 160) for finer temporal resolution (~101 frames vs ~51)
- Mixup augmentation (alpha=0.2) in training loop for better generalisation
- Output dir: speech_cadlif_dyned_output_v4
Carried from v3:
- Quantised magnitudes (continuous) with boost option
- cAdLIF neurons with learnable delays
- OneCycleLR scheduler
- AdamW optimiser
Architecture based on:
- Hammouamri et al. (2024) "Co-learning synaptic delays, weights and adaptation"
- Bittar & Garner (2022) "A surrogate gradient spiking baseline for speech command recognition"
"""
import argparse
import math
import os
import sys
import time
import zipfile
import matplotlib
try:
matplotlib.use("Agg")
except Exception:
pass
import matplotlib.pyplot as plt
import numpy as np
import torch
torch.backends.nnpack.enabled = False
import torch.nn as nn
import torch.nn.functional as F
from pathlib import Path
from torch.utils.data import DataLoader
from torchaudio.datasets import SPEECHCOMMANDS
from torchaudio.transforms import Resample
import torchaudio.transforms as T
try:
from noisereduce.torchgate import TorchGate
HAS_TORCHGATE = True
except ImportError:
HAS_TORCHGATE = False
print("Warning: noisereduce not installed - skipping spectral gate cleaning")
try:
_script_dir = os.path.dirname(os.path.abspath(__file__))
except NameError:
_script_dir = os.getcwd()
sys.path.insert(0, os.path.join(_script_dir, ".."))
from dyned import dyned_quantise_2d
from vis_utils import dump_plot_data
try:
from sklearn.manifold import TSNE
HAS_TSNE = True
except ImportError:
HAS_TSNE = False
try:
_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
except NameError:
_SCRIPT_DIR = os.getcwd()
OUTPUT_DIR = os.path.join(_SCRIPT_DIR, "speech_cadlif_dyned_output_v4")
os.makedirs(OUTPUT_DIR, exist_ok=True)
SPEECH_LABELS = [
"backward",
"bed",
"bird",
"cat",
"dog",
"down",
"eight",
"five",
"follow",
"forward",
"four",
"go",
"happy",
"house",
"learn",
"left",
"marvin",
"nine",
"no",
"off",
"on",
"one",
"right",
"seven",
"sheila",
"six",
"stop",
"three",
"tree",
"two",
"up",
"visual",
"wow",
"yes",
"zero",
]
# =============================================================================
# DyNED Speech Encoding (mel + MFCC through DyNED quantisation)
# =============================================================================
def make_mel_mfcc_transforms(target_sr=8000, n_mels=80, n_mfcc=80, n_fft=1024, hop_length=80):
"""Create mel and MFCC transforms (call once, reuse across samples)."""
mel_transform = T.MelSpectrogram(
sample_rate=target_sr,
n_mels=n_mels,
n_fft=n_fft,
win_length=200,
hop_length=hop_length,
f_min=20,
f_max=target_sr // 2,
center=True,
pad_mode="reflect",
power=2.0,
normalized=True,
)
mfcc_transform = T.MFCC(
sample_rate=target_sr,
n_mfcc=n_mfcc,
log_mels=True,
melkwargs={
"n_mels": n_mels,
"n_fft": n_fft,
"win_length": 200,
"hop_length": hop_length,
"f_min": 20,
"f_max": target_sr // 2,
},
)
return mel_transform, mfcc_transform
def dyned_encode_waveform(waveform, mel_transform, mfcc_transform, dyned_levels=256, boost=True):
"""Encode waveform via mel-spectrogram + MFCC + DyNED, returning quantised features.
Produces [160, n_time] output: 80 mel + 80 MFCC channels, each DyNED-quantised.
"""
mel_spec = mel_transform(waveform.unsqueeze(0)) # [1, n_mels, T]
mel_spec = (mel_spec + 1e-8).log().squeeze(0) # [n_mels, T]
mfcc = mfcc_transform(waveform.unsqueeze(0)).squeeze(0) # [n_mfcc, T]
# Quantise mel and MFCC independently (different value ranges need separate normalisation)
mel_encoded = dyned_quantise_2d(mel_spec.numpy(), levels=dyned_levels, boost=boost)
mfcc_encoded = dyned_quantise_2d(mfcc.numpy(), levels=dyned_levels, boost=boost)
return torch.cat([mel_encoded, mfcc_encoded], dim=0) # [160, T]
# =============================================================================
# Surrogate Gradient
# =============================================================================
class ATanSurrogate(torch.autograd.Function):
"""Arctangent surrogate gradient for spiking neurons (Fang et al., 2021)."""
@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
# =============================================================================
# Constraint bounds (from Bittar & Garner, 2022; Hammouamri et al., 2024)
ALPHA_MIN = math.exp(-1.0 / 5.0) # 0.8187 - fast membrane leak
ALPHA_MAX = math.exp(-1.0 / 25.0) # 0.9608 - slow membrane leak
BETA_MIN = math.exp(-1.0 / 30.0) # 0.9672 - fast adaptation leak
BETA_MAX = math.exp(-1.0 / 120.0) # 0.9917 - slow adaptation leak
THRESHOLD = 0.5
class cAdLIFNeuron(nn.Module):
"""Constrained Adaptive Leaky Integrate-and-Fire neuron.
Equations:
w[t] = beta * w[t-1] + (1 - beta) * a * u[t-1] + b * s[t-1]
u[t] = alpha * u[t-1] + (1 - alpha) * (I[t] - w[t-1]) - theta * s[t-1]
s[t] = Theta(u[t] - theta)
Learnable per-neuron parameters:
alpha - membrane leak rate, in [0.8187, 0.9608]
beta - adaptation leak rate, in [0.9672, 0.9917]
a - subthreshold adaptation coupling, in [0, 1]
b - spike-triggered adaptation, in [0, 2]
"""
def __init__(self, size):
super().__init__()
self.size = size
# Initialize with uniform spread within valid ranges
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):
"""Single timestep update.
Args:
I_t: input current [batch, size]
u_prev: previous membrane potential [batch, size]
w_prev: previous adaptation current [batch, size]
s_prev: previous spike output [batch, size]
Returns:
s: spikes [batch, size]
u: membrane potential [batch, size]
w: adaptation current [batch, size]
"""
alpha, beta, a, b = self._constrain()
# Adaptation current update
w = beta * w_prev + (1.0 - beta) * a * u_prev + b * s_prev
# Membrane potential update
u = alpha * u_prev + (1.0 - alpha) * (I_t - w_prev) - THRESHOLD * s_prev
# Spike generation
s = spike_fn(u - THRESHOLD)
return s, u, w
# =============================================================================
# Learnable Delays
# =============================================================================
def apply_delays(h_seq, delays, max_delay):
"""Apply per-neuron learnable delays with differentiable interpolation.
Args:
h_seq: [T, B, N] - transformed input sequence
delays: [N] - continuous delay values (will be clamped to [0, max_delay])
max_delay: int - maximum delay in timesteps
Returns:
[T, B, N] - delayed sequence
"""
T, B, N = h_seq.shape
device = h_seq.device
# Pad with zeros at the start (so we can look back max_delay steps)
pad = torch.zeros(max_delay, B, N, device=device, dtype=h_seq.dtype)
h_padded = torch.cat([pad, h_seq], dim=0) # [T + max_delay, B, N]
T_pad = T + max_delay
# Clamp delays - detach for index computation, keep gradient for interpolation
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() # [N], gradient flows through here
# Compute time indices: for output time t, neuron n uses h_padded[t + max_delay - d]
t_range = torch.arange(T, device=device) # [T]
idx_floor = (t_range.unsqueeze(1) + max_delay - d_floor.unsqueeze(0)).clamp(0, T_pad - 1) # [T, N]
idx_ceil = (t_range.unsqueeze(1) + max_delay - d_ceil.unsqueeze(0)).clamp(0, T_pad - 1) # [T, N]
# Expand for batch dim and gather
idx_f = idx_floor.unsqueeze(1).expand(T, B, N) # [T, B, N]
idx_c = idx_ceil.unsqueeze(1).expand(T, B, N) # [T, B, N]
val_floor = torch.gather(h_padded, 0, idx_f) # [T, B, N]
val_ceil = torch.gather(h_padded, 0, idx_c) # [T, B, N]
# Differentiable interpolation
frac_exp = frac.view(1, 1, N)
return (1.0 - frac_exp) * val_floor + frac_exp * val_ceil
# =============================================================================
# cAdLIF Speech SNN (v3 - quantised magnitude input, larger hidden)
# =============================================================================
class cAdLIFSpeechSNN(nn.Module):
"""SNN with cAdLIF neurons and learnable delays for speech classification.
Architecture:
Input (n_freq x T) -> Linear + LN + Delay + cAdLIF + Drop (layer 1)
-> Linear + LN + Delay + cAdLIF + Drop (layer 2)
-> Linear + LN + Delay + cAdLIF + Drop (layer 3)
-> Linear + LIF readout (no spike, softmax accumulation)
"""
def __init__(self, n_freq_bins=513, hidden_sizes=(2048, 1024, 512), num_outputs=35, max_delay=20, dropout=0.1):
super().__init__()
self.hidden_sizes = hidden_sizes
self.num_outputs = num_outputs
self.max_delay = max_delay
self.num_layers = len(hidden_sizes)
# Build layers dynamically
self.fc_layers = nn.ModuleList()
self.ln_layers = nn.ModuleList()
self.cadlif_layers = nn.ModuleList()
self.drop_layers = nn.ModuleList()
self.delay_params = nn.ParameterList()
in_size = n_freq_bins
for i, h_size in enumerate(hidden_sizes):
self.fc_layers.append(nn.Linear(in_size, h_size))
self.ln_layers.append(nn.LayerNorm(h_size))
self.cadlif_layers.append(cAdLIFNeuron(h_size))
self.delay_params.append(nn.Parameter(torch.zeros(h_size)))
self.drop_layers.append(nn.Dropout(dropout))
in_size = h_size
# Readout: last hidden -> classes (infinite threshold, no spike)
self.fc_out = nn.Linear(hidden_sizes[-1], num_outputs)
self.alpha_out_raw = nn.Parameter(torch.empty(num_outputs).uniform_(ALPHA_MIN, ALPHA_MAX))
# Weight initialization
for fc in self.fc_layers:
nn.init.kaiming_uniform_(fc.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 forward(self, x):
"""
Args:
x: [batch, n_freq_bins, n_time_frames]
Returns:
output: [batch, num_outputs] - accumulated softmax votes
"""
B = x.size(0)
T = x.size(2)
device = x.device
seq = x.permute(2, 0, 1) # [T, B, freq]
# Process through each cAdLIF layer
for i in range(self.num_layers):
h_size = self.hidden_sizes[i]
# Linear + LayerNorm + Delay
h = self.fc_layers[i](seq) # [T, B, h_size]
h = self.ln_layers[i](h)
h = apply_delays(h, self.delay_params[i].clamp(0, self.max_delay), self.max_delay)
# Run cAdLIF
s_list = []
u = torch.zeros(B, h_size, device=device)
w = torch.zeros(B, h_size, device=device)
s = torch.zeros(B, h_size, device=device)
for t in range(T):
s, u, w = self.cadlif_layers[i](h[t], u, w, s)
s_list.append(s)
seq = torch.stack(s_list) # [T, B, h_size]
seq = self.drop_layers[i](seq)
# Readout: LIF with infinite threshold, softmax accumulation
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(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, C]
output = torch.sum(F.softmax(m_out, dim=2), dim=0) # [B, C]
return output
def diagnostic_forward(self, x):
"""Forward pass that returns per-layer spike and membrane data for monitoring."""
B = x.size(0)
T = x.size(2)
device = x.device
seq = x.permute(2, 0, 1)
layer_data = {}
for i in range(self.num_layers):
h_size = self.hidden_sizes[i]
h = self.fc_layers[i](seq)
h = self.ln_layers[i](h)
h = apply_delays(h, self.delay_params[i].clamp(0, self.max_delay), self.max_delay)
s_list = []
u_list = []
u = w = s = torch.zeros(B, h_size, device=device)
for t in range(T):
s, u, w = self.cadlif_layers[i](h[t], u, w, s)
s_list.append(s)
u_list.append(u)
seq = torch.stack(s_list)
layer_data[f"spk{i + 1}"] = seq
layer_data[f"mem{i + 1}"] = torch.stack(u_list)
# Readout
alpha_out = self._get_alpha_out()
u_out = torch.zeros(B, self.num_outputs, device=device)
mem_list = []
for t in range(T):
cur = self.fc_out(seq[t])
u_out = alpha_out * u_out + (1.0 - alpha_out) * cur
mem_list.append(u_out)
layer_data["mem_out"] = torch.stack(mem_list)
return layer_data
# =============================================================================
# Dataset
# =============================================================================
class DyNEDSpeechDataset(torch.utils.data.Dataset):
"""Speech Commands dataset with pre-computed DyNED quantised magnitudes cached to disk."""
def __init__(
self,
subset="training",
n_fft=1024,
hop_length=160,
dyned_levels=256,
boost=True,
add_noise=True,
noise_level=0.005,
target_sr=8000,
cache_dir="../assets",
):
self.n_fft = n_fft
self.hop_length = hop_length
self.dyned_levels = dyned_levels
self.boost = boost
self.is_training = subset == "training"
self.add_noise = add_noise and self.is_training
self.noise_level = noise_level
self.target_sr = target_sr
self.labels = SPEECH_LABELS
self.label_to_idx = {label: idx for idx, label in enumerate(self.labels)}
boost_tag = "_boost" if boost else "_quant"
cache_path = (
Path(cache_dir)
/ "speech_cadlif_dyned_snn_cached-4"
/ f"dyned_mel_cache_{subset}_sr{target_sr}_nfft{n_fft}_hop{hop_length}_lvl{dyned_levels}{boost_tag}.pt"
)
if cache_path.exists():
print(f"Loading cached data from {cache_path}...")
try:
cache = torch.load(cache_path, weights_only=True)
self.data = cache["data"]
self.label_indices = cache["labels"]
print(f"Loaded {len(self.data)} samples ({self.data.shape})")
return
except (RuntimeError, EOFError, zipfile.BadZipFile) as e:
print(f"Corrupted cache file, rebuilding: {e}")
cache_path.unlink(missing_ok=True)
print(f"Cache not found - pre-computing DyNED data for {subset}...")
self._build_cache(subset, cache_path, cache_dir)
def _build_cache(self, subset, cache_path, cache_dir):
cache_path.parent.mkdir(parents=True, exist_ok=True)
raw_dataset = SPEECHCOMMANDS(cache_dir, download=True, subset=subset)
resampler = Resample(orig_freq=16000, new_freq=self.target_sr)
torchgate = TorchGate(sr=self.target_sr, nonstationary=False) if HAS_TORCHGATE else None
mel_transform, mfcc_transform = make_mel_mfcc_transforms(
target_sr=self.target_sr, n_fft=self.n_fft, hop_length=self.hop_length
)
all_data = []
all_labels = []
total = len(raw_dataset)
for i in range(total):
waveform, sample_rate, label, _, _ = raw_dataset[i]
if waveform.shape[-1] < 16000:
waveform = F.pad(waveform, (0, 16000 - waveform.shape[-1]))
elif waveform.shape[-1] > 16000:
waveform = waveform[..., :16000]
waveform = resampler(waveform)
if torchgate is not None:
with torch.no_grad():
waveform = torchgate(waveform)
if waveform.shape[-1] < self.target_sr:
waveform = F.pad(waveform, (0, self.target_sr - waveform.shape[-1]))
elif waveform.shape[-1] > self.target_sr:
waveform = waveform[..., : self.target_sr]
mean = waveform.mean()
std = waveform.std()
waveform = (waveform - mean) / (std + 1e-8)
wf_pad = F.pad(waveform, (1, 0))
waveform = wf_pad[..., 1:] - 0.97 * wf_pad[..., :-1]
encoded = dyned_encode_waveform(
waveform.squeeze(0),
mel_transform,
mfcc_transform,
self.dyned_levels,
boost=self.boost,
)
all_data.append(encoded)
all_labels.append(self.label_to_idx[label])
if (i + 1) % 1000 == 0 or (i + 1) == total:
pct = (i + 1) / total * 100
print(f" [{subset}] Encoded {i + 1}/{total} samples ({pct:.1f}%)")
self.data = torch.stack(all_data)
self.label_indices = torch.tensor(all_labels, dtype=torch.long)
torch.save({"data": self.data, "labels": self.label_indices}, cache_path)
size_mb = cache_path.stat().st_size / (1024 * 1024)
print(f"Cached {len(self.data)} samples to {cache_path} ({size_mb:.1f} MB)")
def __len__(self):
return len(self.data)
def __getitem__(self, n):
sample = self.data[n]
label_idx = self.label_indices[n]
if self.is_training:
max_shift = sample.shape[1] // 10
shift = torch.randint(-max_shift, max_shift + 1, (1,)).item()
if shift != 0:
sample = torch.roll(sample, shifts=shift, dims=1)
if torch.rand(1).item() < 0.3:
n_freq = sample.shape[0]
mask_width = torch.randint(1, n_freq // 10 + 1, (1,)).item()
mask_start = torch.randint(0, n_freq - mask_width, (1,)).item()
sample = sample.clone()
sample[mask_start : mask_start + mask_width, :] = 0
if torch.rand(1).item() < 0.3:
n_time = sample.shape[1]
mask_width = torch.randint(1, n_time // 10 + 1, (1,)).item()
mask_start = torch.randint(0, n_time - mask_width, (1,)).item()
sample = sample.clone()
sample[:, mask_start : mask_start + mask_width] = 0
if self.add_noise:
sample = sample.clone()
sample = sample + torch.randn_like(sample) * self.noise_level
return sample, label_idx
# =============================================================================
# Data Setup
# =============================================================================
def setup_dataloaders(
batch_size=512, num_workers=4, n_fft=1024, hop_length=160, dyned_levels=256, target_sr=8000, boost=True
):
train_dataset = DyNEDSpeechDataset(
subset="training",
n_fft=n_fft,
hop_length=hop_length,
dyned_levels=dyned_levels,
boost=boost,
target_sr=target_sr,
)
test_dataset = DyNEDSpeechDataset(
subset="testing",
n_fft=n_fft,
hop_length=hop_length,
dyned_levels=dyned_levels,
boost=boost,
add_noise=False,
target_sr=target_sr,
)
loader_kwargs = dict(pin_memory=True)
if num_workers > 0:
loader_kwargs["persistent_workers"] = False
loader_kwargs["prefetch_factor"] = 4
train_loader = DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
**loader_kwargs,
)
test_loader = DataLoader(
test_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_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=0.1,
weight_decay=0.01,
label_smoothing=0.1,
mixup_alpha=0.2,
trial=None,
):
# Separate parameter groups: delays get higher learning rate
delay_params = list(net.delay_params)
delay_ids = {id(p) for p in delay_params}
weight_params = [p for p in net.parameters() if id(p) not in delay_ids]
optimizer = torch.optim.AdamW(
[
{"params": weight_params, "lr": lr, "weight_decay": weight_decay},
{"params": delay_params, "lr": lr_delay, "weight_decay": 0.0},
]
)
# OneCycleLR - same as the 98.86% script
scheduler = torch.optim.lr_scheduler.OneCycleLR(
optimizer,
max_lr=[lr, lr_delay],
total_steps=num_epochs * len(train_loader),
pct_start=0.1,
anneal_strategy="cos",
)
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: {device}")
print(f"Batch size: {train_loader.batch_size}")
if device == "cuda":
print(f"GPU memory: {torch.cuda.memory_allocated() / 1024**3:.2f} GB allocated")
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
for i, (data, targets) in enumerate(train_loader):
data = data.to(device, non_blocking=True)
targets = targets.to(device, non_blocking=True)
# Mixup augmentation
use_mixup = mixup_alpha > 0
if use_mixup:
lam = np.random.beta(mixup_alpha, mixup_alpha)
idx = torch.randperm(data.size(0), device=device)
data = lam * data + (1.0 - lam) * data[idx]
targets_a, targets_b = targets, targets[idx]
optimizer.zero_grad()
if scaler is not None:
with torch.amp.autocast("cuda"):
output = net(data)
if use_mixup:
loss = lam * F.cross_entropy(output, targets_a, label_smoothing=label_smoothing) + (
1.0 - lam
) * F.cross_entropy(output, targets_b, label_smoothing=label_smoothing)
else:
loss = F.cross_entropy(output, targets, label_smoothing=label_smoothing)
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 = net(data)
if use_mixup:
loss = lam * F.cross_entropy(output, targets_a, label_smoothing=label_smoothing) + (
1.0 - lam
) * F.cross_entropy(output, targets_b, label_smoothing=label_smoothing)
else:
loss = F.cross_entropy(output, targets, label_smoothing=label_smoothing)
loss.backward()
torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)
optimizer.step()
scheduler.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][:32].to(device)
layer_data = net.diagnostic_forward(diag_data)
firing_rates = {}
for key in layer_data:
if key.startswith("spk"):
firing_rates[key] = layer_data[key].mean().item()
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"])
# Print neuron and delay diagnostics
alpha_vals = [
net.cadlif_layers[i].alpha_raw.clamp(ALPHA_MIN, ALPHA_MAX).mean().item() for i in range(net.num_layers)
]
delay_vals = [net.delay_params[i].clamp(0, net.max_delay).mean().item() for i in range(net.num_layers)]
fr_vals = [firing_rates.get(f"spk{i + 1}", 0) for i in range(net.num_layers)]
alpha_str = ", ".join(f"{a:.3f}" for a in alpha_vals)
delay_str = ", ".join(f"{d:.1f}" for d in delay_vals)
fr_str = ", ".join(f"{f:.3f}" for f in fr_vals)
print(
f"Epoch {epoch + 1}: Acc = {test_acc:.4f} | Loss = {avg_epoch_loss:.4f} | "
f"LR = {current_lr:.6f} | alpha = [{alpha_str}] | "
f"Delay = [{delay_str}] | FR = [{fr_str}] | {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_cadlif_speech_model.pth"),
)
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 = []
num_classes = len(SPEECH_LABELS)
per_class_correct = np.zeros(num_classes)
per_class_total = np.zeros(num_classes)
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) # [B, C] - accumulated softmax votes
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(num_classes):
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((num_classes, num_classes), 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": {SPEECH_LABELS[i]: float(per_class_acc[i]) for i in range(num_classes)},
"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"])
# Save CSV
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 SPEECH_LABELS)
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 SPEECH_LABELS:
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="")
print(f"Saved metrics to {filepath}")
best_epoch = np.argmax(accuracies)
print(f"\nBest epoch: {best_epoch + 1} - accuracy: {accuracies[best_epoch]:.4f}")
# Loss/accuracy plot
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("cAdLIF Speech SNN Training Progress (v4)")
lines1, labels1 = ax1.get_legend_handles_labels()
lines2, labels2 = ax2.get_legend_handles_labels()
ax2.legend(lines1 + lines2, labels1 + labels2, loc="upper right")
plt.savefig(os.path.join(OUTPUT_DIR, "training_progress.png"), dpi=150)
plt.close()
dump_plot_data(OUTPUT_DIR, "training_progress", epochs=epochs, train_loss=losses, test_accuracy=accuracies)
# Firing rates
if metrics["layer_firing_rates"]:
fr_history = metrics["layer_firing_rates"]
keys = sorted(fr_history[0].keys())
plt.figure(figsize=(12, 6))
fr_per_key = {}
for key in keys:
rates = [fr[key] for fr in fr_history]
fr_per_key[key] = np.array(rates)
plt.plot(epochs, rates, label=key, linewidth=1.5)
plt.xlabel("Epoch")
plt.ylabel("Mean Firing Rate")
plt.title("Per-Layer Firing Rates")
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "firing_rates.png"), dpi=150)
plt.close()
dump_plot_data(OUTPUT_DIR, "firing_rates", epochs=epochs, keys=np.array(keys), **fr_per_key)
# LR schedule
plt.figure(figsize=(10, 4))
plt.plot(epochs, metrics["learning_rate"])
plt.xlabel("Epoch")
plt.ylabel("Learning Rate")
plt.title("Learning Rate Schedule")
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "lr_schedule.png"), dpi=150)
plt.close()
dump_plot_data(OUTPUT_DIR, "lr_schedule", epochs=epochs, learning_rate=np.array(metrics["learning_rate"]))
def visualize_confusion_matrix_plot(cm):
num_classes = len(SPEECH_LABELS)
fig, ax = plt.subplots(figsize=(11, 10), dpi=180)
cm_normalized = cm.astype(float) / (cm.sum(axis=1, keepdims=True) + 1e-8)
im = ax.imshow(cm_normalized, interpolation="nearest", cmap="Blues", vmin=0, vmax=1)
ax.set_xticks(range(num_classes))
ax.set_yticks(range(num_classes))
ax.set_xticklabels(SPEECH_LABELS, rotation=60, ha="right", fontsize=11)
ax.set_yticklabels(SPEECH_LABELS, fontsize=11)
ax.set_xlabel("Predicted", fontsize=14)
ax.set_ylabel("True", fontsize=14)
ax.set_title("Speech Commands confusion matrix (normalised)", fontsize=14, pad=12)
# Only annotate cells whose value is large enough to read at this size.
for i in range(num_classes):
for j in range(num_classes):
v = cm_normalized[i, j]
if v >= 0.10:
ax.text(
j, i, f"{v:.2f}",
ha="center", va="center", fontsize=7,
color="white" if v >= 0.5 else "black",
)
cbar = plt.colorbar(im, ax=ax, fraction=0.04, pad=0.02)
cbar.ax.tick_params(labelsize=10)
cbar.set_label("Normalised count", fontsize=11)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "confusion_matrix.png"), bbox_inches="tight", facecolor="white")
plt.close()
dump_plot_data(OUTPUT_DIR, "confusion_matrix", cm=cm, cm_normalized=cm_normalized, labels=np.array(SPEECH_LABELS))
np.savetxt(
os.path.join(OUTPUT_DIR, "confusion_matrix.csv"),
cm,
delimiter=",",
fmt="%d",
header=",".join(SPEECH_LABELS),
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=(14, 12))
scatter = plt.scatter(embedded[:, 0], embedded[:, 1], c=targets, cmap="nipy_spectral", alpha=0.6, s=5)
cbar = plt.colorbar(scatter, ticks=range(len(SPEECH_LABELS)))
cbar.set_ticklabels(SPEECH_LABELS)
cbar.ax.tick_params(labelsize=6)
plt.title("t-SNE of Output Representations")
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "tsne.png"), dpi=150)
plt.close()
dump_plot_data(OUTPUT_DIR, "tsne", embedded=embedded, targets=targets, labels=np.array(SPEECH_LABELS))
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=(14, 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=8)
plt.ylabel("Accuracy (%)")
plt.title("Per-Class Accuracy on Test Set")
plt.ylim(0, 100)
plt.grid(True, axis="y", alpha=0.3)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "per_class_accuracy.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR,
"per_class_accuracy",
labels_sorted=np.array(labels_sorted),
accs_sorted=np.asarray(accs_sorted),
labels=np.array(labels),
accs=np.asarray(accs),
)
print("Saved per-class accuracy plot")
def collect_diagnostic_batches(net, test_loader, device, max_samples=4000):
"""Run diagnostic_forward across multiple test batches, concatenating the outputs.
Returns a single dict matching diagnostic_forward's structure plus the
accumulated input data and targets. Used for per-class statistics that
need representative coverage across the 35 classes.
"""
net.eval()
accumulated = {}
inputs = []
targets = []
n = 0
with torch.no_grad():
for data, tgt in test_loader:
data = data.to(device)
ld = net.diagnostic_forward(data)
for k, v in ld.items():
accumulated.setdefault(k, []).append(v.cpu())
inputs.append(data.cpu())
targets.append(tgt)
n += data.size(0)
if n >= max_samples:
break
out = {k: torch.cat(v, dim=1) for k, v in accumulated.items()} # cat over batch dim (T, B, ...)
return out, torch.cat(inputs, dim=0), torch.cat(targets, dim=0)
def visualize_network_activity(input_data, layer_data):
"""Four-panel inference snapshot for sample 0: input raster, hidden L3
raster, hidden L3 firing-rate distribution, output membrane potentials."""
spk_l3 = layer_data["spk3"] # [T, B, hidden]
mem_out = layer_data["mem_out"] # [T, B, num_classes]
fig, axes = plt.subplots(2, 2, figsize=(15, 12))
axes[0, 0].imshow(input_data[0].cpu().numpy(), aspect="auto", origin="lower", cmap="binary")
axes[0, 0].set_title("Input Spike Train (Sample 0)")
axes[0, 0].set_xlabel("Time Frame")
axes[0, 0].set_ylabel("Frequency Bin")
axes[0, 1].imshow(spk_l3[:, 0].cpu().numpy(), aspect="auto", cmap="binary")
axes[0, 1].set_title("Hidden Layer 3 Spike Raster (Sample 0)")
axes[0, 1].set_xlabel("Neuron Index")
axes[0, 1].set_ylabel("Time Step")
rates = spk_l3.mean(dim=0).cpu().numpy().flatten()
axes[1, 0].hist(rates, bins=50, color="steelblue")
axes[1, 0].set_title("Hidden Layer 3 Firing Rate Distribution")
axes[1, 0].set_xlabel("Firing Rate")
axes[1, 0].set_ylabel("Count")
axes[1, 1].plot(mem_out[:, 0].cpu().numpy())
axes[1, 1].set_title("Output Membrane Potentials per Class (Sample 0)")
axes[1, 1].set_xlabel("Time Step")
axes[1, 1].set_ylabel("Membrane Potential")
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "network_activity.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR,
"network_activity",
input_sample0=input_data[0],
spk_l3_sample0=spk_l3[:, 0],
firing_rates=rates,
mem_out_sample0=mem_out[:, 0],
)
print("Saved network activity")
def visualize_layer_spike_rasters(layer_data):
"""Four-panel raster: spk1, spk2, spk3, mem_out (heatmap) for sample 0."""
fig, axes = plt.subplots(2, 2, figsize=(16, 12))
panels = [
("spk1", "Layer 1 (Hidden)"),
("spk2", "Layer 2 (Hidden)"),
("spk3", "Layer 3 (Hidden)"),
("mem_out", "Output Layer (Membrane Potentials)"),
]
panel_arrs = {}
for ax, (key, title) in zip(axes.flat, panels):
d = layer_data[key][:, 0].cpu().numpy()
panel_arrs[key] = d
cmap = "binary" if key.startswith("spk") else "viridis"
ax.imshow(d, aspect="auto", cmap=cmap, interpolation="nearest")
ax.set_title(title)
ax.set_xlabel("Neuron Index")
ax.set_ylabel("Time Step")
if key.startswith("spk"):
ax.text(
0.02,
0.98,
f"Rate: {d.mean():.3f}",
transform=ax.transAxes,
va="top",
fontsize=9,
color="red",
bbox=dict(boxstyle="round", facecolor="white", alpha=0.8),
)
plt.suptitle("Per-Layer Activity (Sample 0)", fontsize=14)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "layer_spike_rasters.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR,
"layer_spike_rasters",
spk1=panel_arrs["spk1"],
spk2=panel_arrs["spk2"],
spk3=panel_arrs["spk3"],
mem_out=panel_arrs["mem_out"],
)
print("Saved layer spike rasters")
def visualize_membrane_distributions(layer_data):
"""Spike-count histograms for hidden layers (cAdLIF exposes spikes, not
membrane potentials internally) plus the output-layer membrane distribution."""
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
panels = [
("spk1", "Layer 1 Spike Counts per Neuron", "spike"),
("spk2", "Layer 2 Spike Counts per Neuron", "spike"),
("spk3", "Layer 3 Spike Counts per Neuron", "spike"),
("mem_out", "Output Membrane Potentials", "membrane"),
]
panel_data = {}
for ax, (key, title, kind) in zip(axes.flat, panels):
d = layer_data[key].cpu().numpy()
if kind == "spike":
counts = d.sum(axis=0).flatten()
panel_data[key] = counts
ax.hist(counts, bins=50, density=True, color="steelblue", alpha=0.8)
ax.set_xlabel("Spike Count over Trial")
else:
vals = d.flatten()
panel_data[key] = vals
ax.hist(vals, bins=100, density=True, color="darkorange", alpha=0.8)
ax.set_xlabel("Membrane Potential")
ax.set_title(title)
ax.set_ylabel("Density")
ax.text(
0.02,
0.98,
f"Mean: {d.mean():.3f}\nStd: {d.std():.3f}",
transform=ax.transAxes,
va="top",
fontsize=9,
bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.8),
)
plt.suptitle("Activity Distributions", fontsize=14)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "membrane_distributions.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR,
"membrane_distributions",
spk1_counts=panel_data["spk1"],
spk2_counts=panel_data["spk2"],
spk3_counts=panel_data["spk3"],
mem_out_values=panel_data["mem_out"],
)
print("Saved membrane distributions")
def visualize_per_class_spikes(layer_data, targets):
"""Per-class average spike patterns (output layer, binarised at THRESHOLD).
Requires aggregated diagnostic batches so all classes have samples."""
mem_out = layer_data["mem_out"] # [T, B, num_classes]
if isinstance(targets, torch.Tensor):
targets = targets.cpu()
cols = 7
rows = math.ceil(len(SPEECH_LABELS) / cols)
fig, axes = plt.subplots(rows, cols, figsize=(3 * cols, 3 * rows))
T_dim = mem_out.shape[0]
C_dim = mem_out.shape[2]
per_class_arr = np.full((len(SPEECH_LABELS), T_dim, C_dim), np.nan, dtype=np.float32)
for cls_idx, ax in enumerate(axes.flat):
if cls_idx >= len(SPEECH_LABELS):
ax.axis("off")
continue
mask = targets == cls_idx
if mask.sum() == 0:
ax.set_title(SPEECH_LABELS[cls_idx], fontsize=7)
ax.text(0.5, 0.5, "no samples", transform=ax.transAxes, ha="center", va="center", fontsize=7)
continue
class_spikes = (mem_out[:, mask, :] > THRESHOLD).float().mean(dim=1).cpu().numpy()
per_class_arr[cls_idx] = class_spikes
ax.imshow(class_spikes, aspect="auto", cmap="binary", interpolation="nearest")
ax.set_title(SPEECH_LABELS[cls_idx], fontsize=7)
ax.tick_params(labelsize=5)
plt.suptitle("Per-Class Average Spike Patterns (Output Layer)", fontsize=14)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "per_class_spikes.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR,
"per_class_spikes",
per_class_arr=per_class_arr,
class_labels=np.array(SPEECH_LABELS),
grid=np.array([rows, cols]),
)
print("Saved per-class spikes")
def visualize_weight_distributions(net):
"""Trained weight histograms for the four learned linear layers.
Layout: 2x2 grid (3 hidden layers + readout). Each subplot shows the
layer's weight histogram with the in -> out dimensions in the title and a
stats panel (mu, sigma, n) in the corner. Interior tick labels are hidden
so the four distributions share a single set of reference axes.
"""
panels = [(f"fc_layers.{i}", net.fc_layers[i].weight) for i in range(net.num_layers)]
panels.append(("fc_out", net.fc_out.weight))
titles = {
"fc_layers.0": "Hidden layer 1 (fc_layers.0)",
"fc_layers.1": "Hidden layer 2 (fc_layers.1)",
"fc_layers.2": "Hidden layer 3 (fc_layers.2)",
"fc_out": "Readout (fc_out)",
}
weight_payload = {}
name_list = []
shape_list = []
val_list = []
for name, w in panels:
vals = w.detach().cpu().numpy().flatten()
weight_payload[f"weight__{name}"] = vals
name_list.append(name)
shape_list.append(list(w.shape))
val_list.append(vals)
if len(panels) == 4:
fig, axes = plt.subplots(2, 2, figsize=(11, 8), dpi=180, sharex=True, sharey=True)
axes = axes.flatten()
else:
fig, axes = plt.subplots(1, len(panels), figsize=(5 * len(panels), 5), dpi=180, sharex=True, sharey=True)
if len(panels) == 1:
axes = [axes]
all_w = np.concatenate(val_list)
xlim = (np.quantile(all_w, 0.0005), np.quantile(all_w, 0.9995))
for ax, (name, w), vals in zip(axes, panels, val_list):
ax.hist(vals, bins=80, color="#3498DB", edgecolor="white", alpha=0.92, density=True)
ax.axvline(0, color="#7F8C8D", linewidth=0.7, linestyle="--", alpha=0.7)
in_dim, out_dim = int(w.shape[1]), int(w.shape[0]) # PyTorch Linear stores (out, in)
ax.set_title(
f"{titles.get(name, name)} - {in_dim} -> {out_dim}",
fontsize=11,
fontweight="bold",
)
ax.text(
0.97,
0.95,
f"Mean: {vals.mean():+.4f}\nStd: {vals.std():.4f}",
transform=ax.transAxes,
fontsize=9,
ha="right",
va="top",
bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.8),
)
ax.set_xlim(xlim)
ax.grid(True, axis="y", alpha=0.25)
if len(panels) == 4:
for ax in axes[2:]:
ax.set_xlabel("Weight value", fontsize=11)
axes[0].set_ylabel("Density", fontsize=11)
axes[2].set_ylabel("Density", fontsize=11)
else:
for ax in axes:
ax.set_xlabel("Weight value", fontsize=11)
axes[0].set_ylabel("Density", fontsize=11)
fig.suptitle(
"Trained weight distributions - Speech Commands DyNED + cAdLIF SNN",
fontsize=13,
fontweight="bold",
y=0.995,
)
fig.tight_layout(rect=(0, 0, 1, 0.97))
plt.savefig(os.path.join(OUTPUT_DIR, "weight_distributions.png"), bbox_inches="tight", facecolor="white")
plt.close()
dump_plot_data(
OUTPUT_DIR, "weight_distributions", names=np.array(name_list), shapes=np.array(shape_list), **weight_payload
)
print("Saved weight distributions")
# =============================================================================
# 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")
n_fft = 1024
hop_length = 80
target_sr = 8000
n_feat = 160 # 80 mel + 80 MFCC
num_outputs = 35
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.1)
weight_decay = params.get("weight_decay", 0.01)
hidden_sizes = tuple(params.get("hidden_sizes", [2048, 1024, 512]))
max_delay = params.get("max_delay", 20)
dropout = params.get("dropout", 0.1)
batch_size = params.get("batch_size", 256)
dyned_levels = params.get("dyned_levels", 256)
label_smoothing = params.get("label_smoothing", 0.1)
boost = params.get("boost", True)
mixup_alpha = params.get("mixup_alpha", 0.2)
else:
lr = 5e-4
lr_delay = 0.1
weight_decay = 0.01
hidden_sizes = (2048, 1024, 512)
max_delay = 20
dropout = 0.1
batch_size = 256
dyned_levels = 256
label_smoothing = 0.1
boost = True
mixup_alpha = 0.2
num_epochs = 250
num_workers = min(8, os.cpu_count() - 2) if os.cpu_count() > 2 else 0
if _CLI_WORKERS is not None:
num_workers = _CLI_WORKERS
arch_str = " -> ".join(str(h) for h in hidden_sizes)
boost_str = "boosted" if boost else "quantised"
print("=" * 80)
print("DyNED + cAdLIF SNN for Speech Commands (v4)")
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: Waveform -> 8kHz -> Spectral Gate -> Mel+MFCC -> DyNED (lvl={dyned_levels}, {boost_str}) -> cAdLIF SNN"
)
print(f"Input: {n_feat} features (80 mel + 80 MFCC) x ~101 timesteps ({boost_str} magnitudes)")
print(f"Architecture: {n_feat} -> {arch_str} -> {num_outputs}")
print(f"Neurons: cAdLIF (learnable alpha, beta, a, b) + delays (max={max_delay})")
print(f"Batch: {batch_size} | Dropout: {dropout} | Label smoothing: {label_smoothing} | Mixup alpha: {mixup_alpha}")
print(f"LR: {lr} (weights) / {lr_delay} (delays) | Weight decay: {weight_decay}")
print(f"Scheduler: OneCycleLR (10% warmup)")
print(f"Optimizer: AdamW")
print("=" * 80)
net = cAdLIFSpeechSNN(
n_freq_bins=n_feat,
hidden_sizes=hidden_sizes,
num_outputs=num_outputs,
max_delay=max_delay,
dropout=dropout,
).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_dataloaders(
batch_size=batch_size,
num_workers=num_workers,
n_fft=n_fft,
hop_length=hop_length,
dyned_levels=dyned_levels,
target_sr=target_sr,
boost=boost,
)
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,
label_smoothing=label_smoothing,
mixup_alpha=mixup_alpha,
)
print("\n" + "=" * 80)
print("POST-TRAINING ANALYSIS")
print("=" * 80)
analyze_training_metrics(metrics)
final_eval = evaluate(net, test_loader, device, collect_representations=True)
visualize_confusion_matrix_plot(final_eval["confusion_matrix"])
if "representations" in final_eval:
visualize_tsne(final_eval["representations"], final_eval["targets"])
visualize_per_class_accuracy(final_eval["per_class_accuracy"])
# Network dynamics figures from a representative batch + an aggregated
# multi-batch view for per-class statistics that need all 35 classes.
net.eval()
with torch.no_grad():
sample_data, _ = next(iter(test_loader))
sample_data = sample_data.to(device)
sample_layer_data = net.diagnostic_forward(sample_data)
visualize_network_activity(sample_data, sample_layer_data)
visualize_layer_spike_rasters(sample_layer_data)
visualize_membrane_distributions(sample_layer_data)
agg_layer_data, _, agg_targets = collect_diagnostic_batches(
net,
test_loader,
device,
max_samples=4000,
)
visualize_per_class_spikes(agg_layer_data, agg_targets)
visualize_weight_distributions(net)
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}")
# Print learned neuron parameters
print("\nLearned neuron parameters:")
for i in range(net.num_layers):
cadlif = net.cadlif_layers[i]
alpha, beta, a, b = cadlif._constrain()
print(f" Layer {i + 1}: alpha={alpha.mean():.4f} beta={beta.mean():.4f} a={a.mean():.4f} b={b.mean():.4f}")
print(f" Readout alpha: {net._get_alpha_out().mean():.4f}")
print("\nLearned delays:")
for i in range(net.num_layers):
d = net.delay_params[i].clamp(0, net.max_delay)
print(f" Layer {i + 1}: 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,
"n_fft": n_fft,
"hidden_sizes": hidden_sizes,
"boost": boost,
"n_feat": n_feat,
"hop_length": hop_length,
"mixup_alpha": mixup_alpha,
"metrics": {
"train_losses": metrics["train_loss"],
"test_accuracies": metrics["test_accuracy"],
"final_accuracy": metrics["test_accuracy"][-1],
},
},
os.path.join(OUTPUT_DIR, "final_model.pth"),
)
print(f"\nDone! Accuracy: {metrics['test_accuracy'][-1]:.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"))
def optuna_objective(trial):
import optuna
device_str = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device_str)
n_fft = 1024
hop_length = 80
target_sr = 8000
n_feat = 160 # 80 mel + 80 MFCC
num_outputs = 35
lr = trial.suggest_float("lr", 1e-4, 1e-3, log=True)
lr_delay = trial.suggest_float("lr_delay", 1e-2, 1.0, log=True)
weight_decay = trial.suggest_float("weight_decay", 1e-3, 1e-1, log=True)
h1 = trial.suggest_categorical("hidden_1", [1024, 2048])
h2 = trial.suggest_categorical("hidden_2", [512, 1024])
h3 = trial.suggest_categorical("hidden_3", [256, 512])
hidden_sizes = (h1, h2, h3)
max_delay = trial.suggest_categorical("max_delay", [10, 15, 20, 25])
dropout = trial.suggest_float("dropout", 0.05, 0.2)
batch_size = trial.suggest_categorical("batch_size", [128, 256])
dyned_levels = trial.suggest_categorical("dyned_levels", [128, 256])
label_smoothing = trial.suggest_float("label_smoothing", 0.0, 0.2)
boost = trial.suggest_categorical("boost", [True, False])
mixup_alpha = trial.suggest_float("mixup_alpha", 0.0, 0.4)
num_epochs = 50
num_workers = 0 # Avoid shared memory exhaustion with parallel Optuna jobs
if _CLI_WORKERS is not None:
num_workers = _CLI_WORKERS
net = cAdLIFSpeechSNN(
n_freq_bins=n_feat,
hidden_sizes=hidden_sizes,
num_outputs=num_outputs,
max_delay=max_delay,
dropout=dropout,
).to(device)
train_loader, test_loader = setup_dataloaders(
batch_size=batch_size,
num_workers=num_workers,
n_fft=n_fft,
hop_length=hop_length,
dyned_levels=dyned_levels,
target_sr=target_sr,
boost=boost,
)
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,
label_smoothing=label_smoothing,
mixup_alpha=mixup_alpha,
trial=trial,
)
except optuna.TrialPruned:
raise
finally:
del net, train_loader, test_loader
if device_str == "cuda":
torch.cuda.empty_cache()
return max(metrics["test_accuracy"])
def optuna_optimize(n_trials=50, study_name="speech_cadlif_dyned_snn_v4", 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 accuracy: {study.best_trial.value:.4f}")
for key, value in study.best_trial.params.items():
print(f" {key}: {value}")
import json
with open(os.path.join(OUTPUT_DIR, f"{study_name}_best_params.json"), "w") as f:
json.dump({"accuracy": study.best_trial.value, "optuna_epochs": 50, **study.best_trial.params}, f, indent=2)
return study
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Speech Commands cAdLIF DyNED SNN (v4)")
parser.add_argument("--optuna", action="store_true")
parser.add_argument("--n-trials", type=int, default=50)
parser.add_argument("--study-name", type=str, default="speech_cadlif_dyned_snn_v4")
parser.add_argument("--n-jobs", type=int, default=1)
parser.add_argument("--json", type=str, default=None)
parser.add_argument("--workers", type=int, default=None, help="DataLoader num_workers (default: auto)")
parser.add_argument("--gen-data", action="store_true", help="Pre-generate all cache files, then exit")
parser.add_argument("--cpu", action="store_true", help="Force CPU (disable CUDA)")
args = parser.parse_args()
if args.cpu:
os.environ["CUDA_VISIBLE_DEVICES"] = ""
_CLI_WORKERS = args.workers
if args.gen_data:
import multiprocessing
def _build_with_prefix(prefix, cls, kwargs):
"""Build cache with prefixed output lines."""
import sys
_orig_write = sys.stdout.write
_orig_flush = sys.stdout.flush
def _prefixed_write(s):
if s.strip():
_orig_write(f" [{prefix}] {s}")
else:
_orig_write(s)
_orig_flush()
sys.stdout.write = _prefixed_write
try:
cls(**kwargs)
finally:
sys.stdout.write = _orig_write
combos = []
for levels in [128, 256]:
for boost in [True, False]:
btag = "boost" if boost else "quant"
for subset in ["training", "testing"]:
stag = "train" if subset == "training" else "test"
prefix = f"lvl={levels},{btag},{stag}"
combos.append(
(
prefix,
dict(
subset=subset,
n_fft=1024,
hop_length=80,
dyned_levels=levels,
boost=boost,
target_sr=8000,
),
)
)
print(f"Pre-generating {len(combos)} cache files...")
# Ensure raw data is downloaded before spawning parallel workers
# (one archive for all subsets - subset only filters which wavs to list)
SPEECHCOMMANDS("../assets", download=True, subset="training")
for prefix, kwargs in combos:
_build_with_prefix(prefix, DyNEDSpeechDataset, kwargs)
print("Cache generation complete.")
sys.exit(0)
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)