speech_cadlif_dyned_dynedc_snn_cached-8.py
Speech Commands DyNED + DyNEDc SNN - 2,186 lines.
View on GitHub (speech/speech_cadlif_dyned_dynedc_snn_cached-8.py).
Source
Section titled “Source”"""
Speech Commands SNN with DyNED + DyNEDc (ON/OFF binary direct) - cAdLIF + Learnable Delays (v8)
Pipeline (end-to-end binary, fully consistent):
Waveform -> Mel+MFCC -> DyNED quantise -> ON/OFF event channels (binary) [2, 160, 101]
-> DyNEDc compress -> bytes on disk
-> DyNEDc decompress -> ON/OFF binary [2, 160, 101]
-> reshape to [320, 101] -> cAdLIF SNN (binary input, same v4 model)
Why this design:
- DyNEDc IS physically in the data path (compress->decompress every sample).
- The cache file actually shrinks by the compression ratio (real storage win).
- Lossless guarantee holds end-to-end: bit-identical binary in -> out.
- The cAdLIF model trains on the SAME binary representation that DyNEDc operates
on - no train/eval mismatch, no lossy reconstruction hack, no separate "spike
train side metric" framing. One coherent pipeline.
- Standard neuromorphic encoding (ON/OFF events, like DVS cameras and SSC).
Differences from v7 (which had a lossy cumsum reconstruction):
- v7 tried to recover continuous values from ON/OFF via cumsum; this was
severely lossy (~98% clip rate on real DyNED data) because consecutive
quantised levels can change by more than 1 step.
- v8 skips reconstruction entirely. The model sees the ON/OFF binary spike
train directly. cAdLIF is a spiking architecture and natively handles binary
input (input layer is just `Linear(320, hidden)` over binary).
Expected accuracy:
- Lower than v4 DyNED's 92.79% (which uses continuous DyNED-quantised values).
- Comparable to RadLIF / d-cAdLIF results on binary spike inputs (~70-85%).
- Honest result of the unified DyNED->DyNEDc->SNN pipeline.
Carried from v4 DyNED:
- Mel + MFCC features (160 channels x ~101 timesteps)
- cAdLIF neurons with learnable delays
- OneCycleLR scheduler, AdamW optimiser
- Mixup augmentation (alpha=0.2)
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, DyNEDcCompressorV4
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_dynedc_output_v8")
# Module-level CLI overrides (set in __main__ block; default None means "use script defaults")
_CLI_WORKERS = None
_CLI_SUBSET_SIZE = None
_CLI_QUICK_EPOCHS = None
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]
# =============================================================================
# ON/OFF Event Encoding Helpers - DyNEDc lossless round-trip via differencing
# =============================================================================
def encode_on_off(quantised_2d, levels):
"""Encode DyNED-quantised continuous values as ON/OFF event channels.
DyNED's sigma-delta runs PER TIMESTEP across the channel axis (each frame
[n_channels, t] is sigma-delta encoded along the channel direction). So
consecutive channels at a fixed timestep have small deltas - that is the
natural axis for differencing.
Encoding (per timestep t):
1. Per-frame normalisation: idx[c, t] = round((q[c,t] - col_min[t]) /
(col_range[t]) * (levels-1)).
2. Compute Delta[c, t] = idx[c, t] - idx[c-1, t] along the channel axis.
3. Clamp Delta to {-1, 0, +1}; record clip rate.
4. ON: Delta == +1, OFF: Delta == -1.
Args:
quantised_2d: tensor or array [n_channels, n_time], from dyned_quantise_2d
levels: number of DyNED quantisation levels (e.g. 256)
Returns:
on_off: uint8 array [2, n_channels, n_time], values in {0, 1}
axis 0: 0 = ON channel, 1 = OFF channel
meta: dict with reconstruction metadata:
- "starting_idx": uint16 array [n_time] (initial level per timestep, channel 0)
- "col_min": float32 array [n_time]
- "col_max": float32 array [n_time]
- "n_clipped": int (number of cells where |Delta| > 1 was clipped)
"""
arr = quantised_2d.detach().cpu().numpy() if hasattr(quantised_2d, "detach") else np.asarray(quantised_2d)
n_rows, n_time = arr.shape
col_min = arr.min(axis=0).astype(np.float32) # [n_time]
col_max = arr.max(axis=0).astype(np.float32)
col_range = col_max - col_min
col_range_safe = np.where(col_range < 1e-8, 1.0, col_range)
norm = (arr - col_min[None, :]) / col_range_safe[None, :]
idx = np.round(norm * (levels - 1)).clip(0, levels - 1).astype(np.int32)
starting_idx = idx[0, :].astype(np.uint16) # initial level per timestep
delta = idx[1:, :] - idx[:-1, :] # signed integer [n_rows-1, n_time]
clipped_delta = np.clip(delta, -1, 1).astype(np.int8)
n_clipped = int(np.sum(np.abs(delta) > 1))
on = (clipped_delta > 0).astype(np.uint8) # [n_rows-1, n_time]
off = (clipped_delta < 0).astype(np.uint8)
# Pad a leading zero row so shape matches the original n_rows
on = np.concatenate([np.zeros((1, n_time), dtype=np.uint8), on], axis=0)
off = np.concatenate([np.zeros((1, n_time), dtype=np.uint8), off], axis=0)
on_off = np.stack([on, off], axis=0) # [2, n_rows, n_time]
meta = {
"starting_idx": starting_idx,
"col_min": col_min,
"col_max": col_max,
"n_clipped": n_clipped,
}
return on_off, meta
def decode_on_off(on_off, meta, levels):
"""Reconstruct continuous quantised values from ON/OFF event channels.
Inverse of `encode_on_off`. Bit-exact when no large jumps were clipped at
encoding; small reconstruction error otherwise.
Args:
on_off: uint8 array [2, n_channels, n_time] - ON in axis 0, OFF in axis 1
meta: dict returned by encode_on_off (starting_idx, col_min, col_max)
levels: number of DyNED quantisation levels
Returns:
recon: float32 tensor [n_channels, n_time]
"""
on = on_off[0].astype(np.int32)
off = on_off[1].astype(np.int32)
delta = on - off # signed +/-1 or 0 per cell
# First row was a forced zero - restore starting idx as the row-0 cumulative seed.
delta[0, :] = 0
cumulative = np.cumsum(delta, axis=0) # [n_rows, n_time]
idx = cumulative + meta["starting_idx"][None, :].astype(np.int32)
idx = idx.clip(0, levels - 1)
col_range = meta["col_max"] - meta["col_min"]
step = col_range / (levels - 1)
recon = idx.astype(np.float32) * step[None, :] + meta["col_min"][None, :]
return torch.from_numpy(recon)
def dynedc_compress_binary(binary_2d, chunk_size=4):
"""Compress a binary array with DyNEDcCompressorV4.
Returns: (bit_string, info_dict, codec_state). `codec_state` captures the
per-sample mode/Huffman codes/alt-start that the decoder needs (V4's
decompress is stateful per encode call).
"""
flat = binary_2d.astype(np.uint8).flatten()
compressor = DyNEDcCompressorV4(chunk_size=chunk_size)
compressed, info = compressor.compress(flat)
codec_state = {
"mode": compressor._mode,
"huff_codes": dict(compressor._huff_codes),
"alt_start": compressor._alt_start,
}
return compressed, info, codec_state
def dynedc_decompress_binary(compressed_str, shape, codec_state, chunk_size=4):
"""Inverse of `dynedc_compress_binary`. Restores the per-sample compressor
state from `codec_state` so V4's stateful decoder works correctly.
"""
compressor = DyNEDcCompressorV4(chunk_size=chunk_size)
compressor._mode = codec_state["mode"]
compressor._huff_codes = codec_state["huff_codes"]
compressor._alt_start = codec_state["alt_start"]
decompressed_str = compressor.decompress(compressed_str)
flat = np.frombuffer(decompressed_str.encode("ascii"), dtype=np.uint8) - ord("0")
n_expected = int(np.prod(shape))
flat = flat[:n_expected]
return flat.reshape(shape).astype(np.uint8)
# =============================================================================
# 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 (v4 + DyNEDc - quantised magnitude input, larger hidden)
# =============================================================================
class cAdLIFSpeechSNN(nn.Module):
"""SNN with cAdLIF neurons, learnable delays, and DyNEDc compression for speech classification.
Architecture:
Input (n_freq x T) -> DyNEDc compress/decompress (eval only)
-> 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)
# NB: DyNEDc compression happens in the dataset/dataloader (cache stores
# compressed bitstreams; decompression + DyNED reconstruction in __getitem__).
# The model sees continuous reconstructed values identical to v4 DyNED's input.
# 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 DyNEDc-compressed ON/OFF event cache.
Cache format (single .pt file per subset):
{
"compressed": list[bytes] # bit-packed DyNEDc-compressed ON/OFF payloads
"comp_lengths": int64 tensor # uncompressed bit count of each payload (for decode)
"labels": int64 tensor # class indices
"starting_idx": int64 tensor # [N, n_channels] initial level per channel per sample
"col_min": float32 tensor [N, n_channels]
"col_max": float32 tensor [N, n_channels]
"binary_shape": tuple # (2, n_channels, n_time) of the ON/OFF tensor
"stats": dict # mean CR, space saving, mode distribution
}
"""
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",
dynedc_chunk_size=4,
max_samples=None,
):
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.dynedc_chunk_size = dynedc_chunk_size
self.max_samples = max_samples
self.dynedc_stats = None
self._compressor = DyNEDcCompressorV4(chunk_size=dynedc_chunk_size)
self.labels = SPEECH_LABELS
self.label_to_idx = {label: idx for idx, label in enumerate(self.labels)}
boost_tag = "_boost" if boost else "_quant"
subset_tag = f"_strat{max_samples}" if max_samples is not None else ""
cache_filename = (
f"dynedc_onoff_cache_{subset}_sr{target_sr}_nfft{n_fft}"
f"_hop{hop_length}_lvl{dyned_levels}{boost_tag}_chunk{dynedc_chunk_size}{subset_tag}.pt"
)
# Cache lives in the v8 subdir. v7's existing files use the same on-disk
# format (same ON/OFF compressed bytes + codec state + col_min/col_max
# metadata), so we transparently fall back to v7's caches if the v8 file
# doesn't exist yet - no need to rebuild.
cache_path = Path(cache_dir) / "speech_cadlif_dyned_dynedc_snn_cached-8" / cache_filename
sibling_v7 = Path(cache_dir) / "speech_cadlif_dyned_dynedc_snn_cached-7" / cache_filename
if not cache_path.exists() and sibling_v7.exists():
cache_path = sibling_v7
self._cache_path = cache_path
if cache_path.exists():
print(f"Loading ON/OFF compressed cache from {cache_path}...")
try:
cache = torch.load(cache_path, weights_only=False)
self._compressed = cache["compressed"]
self._comp_lengths = cache["comp_lengths"]
self._codec_states = cache["codec_states"]
self.label_indices = cache["labels"]
self._starting_idx = cache["starting_idx"]
self._col_min = cache["col_min"]
self._col_max = cache["col_max"]
self._binary_shape = cache["binary_shape"]
self.dynedc_stats = cache["stats"]
size_mb = cache_path.stat().st_size / (1024 * 1024)
print(
f"Loaded {len(self._compressed)} ON/OFF samples "
f"(file: {size_mb:.1f} MB on disk; "
f"CR={self.dynedc_stats['mean_compression_ratio']:.4f}, "
f"{self.dynedc_stats['mean_space_saving_pct']:.1f}% saved)"
)
return
except (RuntimeError, EOFError, zipfile.BadZipFile, KeyError) as e:
print(f"Corrupted cache file, rebuilding: {e}")
cache_path.unlink(missing_ok=True)
print(f"Cache not found - encoding + ON/OFF compressing for {subset}...")
self._build_cache(subset, cache_path, cache_dir)
@staticmethod
def _bitstr_to_packed(bitstr):
arr = np.frombuffer(bitstr.encode("ascii"), dtype=np.uint8) - ord("0")
return np.packbits(arr).tobytes(), len(arr)
@staticmethod
def _packed_to_bitstr(packed_bytes, n_bits):
arr = np.unpackbits(np.frombuffer(packed_bytes, dtype=np.uint8))[:n_bits]
return (arr + ord("0")).astype(np.uint8).tobytes().decode("ascii")
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
)
compressed_payloads = []
comp_lengths = []
codec_states = []
labels = []
starting_idx_list = []
col_min_list = []
col_max_list = []
ratios = []
modes = []
n_clipped_total = 0
binary_shape = None
from collections import Counter
n_total = len(raw_dataset)
if self.max_samples is not None:
# Stratified subset: roughly equal samples per class.
# Read labels from filenames (raw_dataset._walker) - much faster than
# loading each .wav via raw_dataset[i].
per_class_quota = max(1, self.max_samples // len(self.labels))
class_counts = {lbl: 0 for lbl in self.labels}
indices = []
print(f" [{subset}] scanning {n_total} samples for stratified subset ({per_class_quota}/class)...")
for i, path in enumerate(raw_dataset._walker):
label = os.path.basename(os.path.dirname(path))
if label in class_counts and class_counts[label] < per_class_quota:
indices.append(i)
class_counts[label] += 1
if sum(class_counts.values()) >= self.max_samples:
break
print(
f" [{subset}] selected {len(indices)} samples across {sum(1 for c in class_counts.values() if c > 0)} classes"
)
else:
indices = range(n_total)
total = len(indices) if hasattr(indices, "__len__") else n_total
for loop_i, i in enumerate(indices):
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,
) # continuous quantised [160, T]
on_off, meta = encode_on_off(encoded, self.dyned_levels) # [2, 160, T]
if binary_shape is None:
binary_shape = on_off.shape
compressed_str, info, codec_state = dynedc_compress_binary(on_off, self.dynedc_chunk_size)
packed_bytes, bit_len = self._bitstr_to_packed(compressed_str)
compressed_payloads.append(packed_bytes)
comp_lengths.append(bit_len)
codec_states.append(codec_state)
labels.append(self.label_to_idx[label])
starting_idx_list.append(meta["starting_idx"])
col_min_list.append(meta["col_min"])
col_max_list.append(meta["col_max"])
ratios.append(info["compression_ratio"])
modes.append(info.get("mode", "unknown"))
n_clipped_total += meta["n_clipped"]
if (loop_i + 1) % 1000 == 0 or (loop_i + 1) == total:
pct = (loop_i + 1) / total * 100
print(
f" [{subset}] Encoded+ON/OFF compressed {loop_i + 1}/{total} ({pct:.1f}%, "
f"mean CR so far: {np.mean(ratios):.4f}, "
f"clipped Delta>1 cells: {n_clipped_total})"
)
self._compressed = compressed_payloads
self._comp_lengths = torch.tensor(comp_lengths, dtype=torch.int64)
self._codec_states = codec_states
self.label_indices = torch.tensor(labels, dtype=torch.long)
self._starting_idx = torch.tensor(np.stack(starting_idx_list, axis=0), dtype=torch.int64)
self._col_min = torch.tensor(np.stack(col_min_list, axis=0), dtype=torch.float32)
self._col_max = torch.tensor(np.stack(col_max_list, axis=0), dtype=torch.float32)
self._binary_shape = binary_shape
raw_total_bits = int(np.prod(binary_shape)) * len(labels)
compressed_total_bits = int(self._comp_lengths.sum().item())
n_total_cells = len(labels) * binary_shape[1] * (binary_shape[2] - 1) # delta cells
self.dynedc_stats = {
"mean_compression_ratio": float(np.mean(ratios)),
"mean_space_saving_pct": float((1.0 - np.mean(ratios)) * 100),
"min_ratio": float(np.min(ratios)),
"max_ratio": float(np.max(ratios)),
"n_samples": len(ratios),
"chunk_size": self.dynedc_chunk_size,
"mode_distribution": dict(Counter(modes)),
"raw_total_bits": raw_total_bits,
"compressed_total_bits": compressed_total_bits,
"binary_shape": list(binary_shape),
"n_clipped_cells": n_clipped_total,
"clip_rate_pct": float(100.0 * n_clipped_total / n_total_cells) if n_total_cells > 0 else 0.0,
}
torch.save(
{
"compressed": self._compressed,
"comp_lengths": self._comp_lengths,
"codec_states": self._codec_states,
"labels": self.label_indices,
"starting_idx": self._starting_idx,
"col_min": self._col_min,
"col_max": self._col_max,
"binary_shape": self._binary_shape,
"stats": self.dynedc_stats,
},
cache_path,
)
size_mb = cache_path.stat().st_size / (1024 * 1024)
print(
f"Cached {len(self._compressed)} ON/OFF samples to {cache_path} "
f"({size_mb:.1f} MB on disk; CR={self.dynedc_stats['mean_compression_ratio']:.4f}, "
f"{self.dynedc_stats['mean_space_saving_pct']:.1f}% saved; "
f"clip rate {self.dynedc_stats['clip_rate_pct']:.2f}%)"
)
def __len__(self):
return len(self._compressed)
def __getitem__(self, n):
# --- DyNEDc decompress: bit-packed bytes -> bit string -> ON/OFF binary ---
bitstr = self._packed_to_bitstr(self._compressed[n], int(self._comp_lengths[n].item()))
on_off = dynedc_decompress_binary(
bitstr,
self._binary_shape,
self._codec_states[n],
self.dynedc_chunk_size,
) # uint8 [2, 160, n_time], values in {0, 1}
# --- Feed binary spike train directly to the model.
# Reshape ON/OFF channels into the feature axis: [2, 160, T] -> [320, T].
# Channels 0..159 = ON spikes (level increased between channels)
# Channels 160..319 = OFF spikes (level decreased)
on_off_t = torch.from_numpy(on_off.astype(np.float32)) # [2, 160, T]
sample = on_off_t.reshape(2 * on_off.shape[1], on_off.shape[2]) # [320, T]
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,
dynedc_chunk_size=4,
max_samples=None,
):
train_dataset = DyNEDSpeechDataset(
subset="training",
n_fft=n_fft,
hop_length=hop_length,
dyned_levels=dyned_levels,
boost=boost,
target_sr=target_sr,
dynedc_chunk_size=dynedc_chunk_size,
max_samples=max_samples,
)
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,
dynedc_chunk_size=dynedc_chunk_size,
max_samples=max_samples,
)
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_dynedc_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()
# NB: DyNEDc compression stats are computed once at cache-build time and
# saved to `dynedc_compression_stats.json` next to the cache file.
# See DyNEDSpeechDataset._build_cache and the printout at the end of main().
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("Speech Commands cAdLIF+DyNED+DyNEDc 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()
dump_plot_data(
OUTPUT_DIR,
"training_progress",
epochs=epochs,
losses=losses,
accuracies=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))
rates_matrix = []
for key in keys:
rates = [fr[key] for fr in fr_history]
rates_matrix.append(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),
rates=np.array(rates_matrix),
)
# 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 35 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)",
}
# Stash the data first so we can dump it and compute a shared x range.
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)
# Use the same 2x2 layout regardless of layer count; if there are more or
# fewer than four panels, fall back to a single row.
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:
# Bottom row gets x label; left column gets y label.
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
# ON/OFF event encoding doubles the feature axis: 160 channels -> 320 inputs
# (channels 0-159 = ON spikes, channels 160-319 = OFF spikes)
n_feat = 320
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)
# NB: mixup is disabled by default in v8 because it interpolates between
# binary samples -> produces fractional values that break the binary
# spike-train semantic during training (train/eval distribution mismatch).
mixup_alpha = params.get("mixup_alpha", 0.0)
dynedc_chunk_size = params.get("dynedc_chunk_size", 4)
else:
lr = 2e-3
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.0 # disabled for v8 (binary input - see note above)
dynedc_chunk_size = 4
num_epochs = 250
if _CLI_QUICK_EPOCHS is not None:
num_epochs = _CLI_QUICK_EPOCHS
print(f"** Quick mode: overriding num_epochs to {num_epochs}")
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
max_samples = _CLI_SUBSET_SIZE
if max_samples is not None:
print(f"** Subset mode: limiting train/test caches to first {max_samples} samples")
arch_str = " -> ".join(str(h) for h in hidden_sizes)
boost_str = "boosted" if boost else "quantised"
print("=" * 80)
print("DyNED + DyNEDc (binary in path) + cAdLIF SNN for Speech Commands (v8)")
print("=" * 80)
print(f"Device: {device_str.upper()}")
if torch.cuda.is_available():
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(
f"Storage: Waveform -> Mel+MFCC -> DyNED (lvl={dyned_levels}, {boost_str}) -> ON/OFF events"
f" -> DyNEDc (chunk={dynedc_chunk_size}) compressed bytes on disk"
)
print(f"Runtime: cached bytes -> DyNEDc decompress -> ON/OFF binary [320, 101] -> cAdLIF SNN (binary input)")
print(f"Input: {n_feat} binary features (160 ON channels + 160 OFF channels) x ~101 timesteps")
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"DyNEDc: lossless binary in path (cache shrinks by CR; expected acc 70-85%, lower than v4 DyNED's 92.79%)")
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,
dynedc_chunk_size=dynedc_chunk_size,
max_samples=max_samples,
)
try:
metrics = train_network(
net=net,
train_loader=train_loader,
test_loader=test_loader,
num_epochs=num_epochs,
device=device_str,
lr=lr,
lr_delay=lr_delay,
weight_decay=weight_decay,
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}")
# Report DyNEDc compression stats from the cache
train_stats = getattr(train_loader.dataset, "dynedc_stats", None)
test_stats = getattr(test_loader.dataset, "dynedc_stats", None)
if train_stats:
print(f"\nDyNEDc compression (train cache, chunk_size={train_stats['chunk_size']}):")
print(f" Mean compression ratio: {train_stats['mean_compression_ratio']:.4f}")
print(f" Mean space saving: {train_stats['mean_space_saving_pct']:.2f}%")
print(f" Range: [{train_stats['min_ratio']:.4f}, {train_stats['max_ratio']:.4f}]")
print(f" n_samples: {train_stats['n_samples']}")
print(f" Mode distribution: {train_stats['mode_distribution']}")
if test_stats:
print(f"\nDyNEDc compression (test cache):")
print(f" Mean compression ratio: {test_stats['mean_compression_ratio']:.4f}")
print(f" Mean space saving: {test_stats['mean_space_saving_pct']:.2f}%")
print(f"\nOutputs: {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
# ON/OFF event encoding doubles the feature axis: 160 channels -> 320 inputs
# (channels 0-159 = ON spikes, channels 160-319 = OFF spikes)
n_feat = 320
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 disabled for binary input (would create fractional values)
mixup_alpha = 0.0
# NB: chunk_size only affects compression ratio, not model accuracy (DyNEDc is lossless)
dynedc_chunk_size = trial.suggest_categorical("dynedc_chunk_size", [2, 4, 8])
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,
dynedc_chunk_size=dynedc_chunk_size,
)
try:
metrics = train_network(
net=net,
train_loader=train_loader,
test_loader=test_loader,
num_epochs=num_epochs,
device=device_str,
lr=lr,
lr_delay=lr_delay,
weight_decay=weight_decay,
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_dynedc_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+DyNEDc 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_dynedc_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)")
parser.add_argument(
"--subset-size", type=int, default=None, help="Limit train/test caches to N samples each for quick smoke tests"
)
parser.add_argument(
"--quick-epochs", type=int, default=None, help="Override num_epochs (use with --subset-size for fast iteration)"
)
args = parser.parse_args()
_CLI_SUBSET_SIZE = args.subset_size
_CLI_QUICK_EPOCHS = args.quick_epochs
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"
# Multiple chunk sizes so Optuna trials don't pay the
# lazy compression-stats computation cost on first encounter.
for chunk_size in [2, 4, 8]:
prefix = f"lvl={levels},{btag},{stag},chunk={chunk_size}"
combos.append(
(
prefix,
dict(
subset=subset,
n_fft=1024,
hop_length=80,
dyned_levels=levels,
boost=boost,
target_sr=8000,
dynedc_chunk_size=chunk_size,
),
)
)
print(f"Pre-generating {len(combos)} cache + stats combinations (8 .pt caches x 3 chunk-size stats 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)