ssc_cadlif_dyned_dynedc_snn_v2.py
SSC DyNED + DyNEDc SNN - 2,156 lines.
View on GitHub (speech-neuro/ssc_cadlif_dyned_dynedc_snn_v2.py).
Source
Section titled “Source”"""
SSC (Spiking Speech Commands) SNN with DyNED Encoding + DyNEDc Compression (v2)
Pipeline (training):
SSC cochlea spikes -> dense histogram [700, T] -> log1p -> DyNED-quantise
(continuous, level-quantised values) -> cAdLIF SNN
Pipeline (compression measurement at inference):
Best model -> forward to cAdLIF spike outputs (binary spike trains, native
to cAdLIF) -> DyNEDc compress per layer per sample -> stats JSON
DyNEDc-in-the-data-path notes:
- DyNED runs at cache build time and produces continuous level-quantised
values that feed the cAdLIF stack directly. We do not binarise between
DyNED and cAdLIF: the previous ON/OFF route created a train/eval mismatch
and capped accuracy.
- The cAdLIF layers themselves emit binary spike trains during inference;
`measure_compression_stats` runs DyNEDc directly on those spike outputs
(`spk1`, `spk2`). This makes the thesis claim "DyNEDc compresses spike
trains" literal: the compressed objects are the actual binary spike trains
the cAdLIF layers emit during inference.
- DyNEDc compression itself is lossless. It runs once over the test set on
the trained model; results are written to dynedc_compression_stats.json.
- Output dir: ssc_cadlif_dyned_dynedc_output_v2
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
try:
import h5py
HAS_H5PY = True
except ImportError:
HAS_H5PY = False
print("Warning: h5py not installed - cannot load SSC dataset")
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 (
sigma_delta_quantisation,
generate_step_signal,
DyNEDcCompressorV4,
dyned_encode_dense,
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, "ssc_cadlif_dyned_dynedc_output_v2")
os.makedirs(OUTPUT_DIR, exist_ok=True)
_CLI_WORKERS = None # Set by --workers flag
# Module-level CLI overrides (set in __main__ block; default None means "use script defaults")
_CLI_SUBSET_SIZE = None
_CLI_QUICK_EPOCHS = None
SSC_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",
]
# Default SSC h5 asset path (relative to script location)
_DEFAULT_SSC_DIR = os.path.join("..", "assets", "SSC")
# Number of cochlea channels in SSC
SSC_N_CHANNELS = 700
# Number of time bins for dense histogram (gives ~20ms per bin over ~1s)
SSC_N_BINS = 50
# =============================================================================
# SSC Dense Histogram + DyNED Encoding
# =============================================================================
def _encode_ssc_sample(args):
"""Worker function for multiprocessing: encode a single SSC sample.
Pipeline: dense histogram -> log1p -> DyNED-quantise (continuous values).
Returns the float32 [n_channels, n_bins] tensor that is fed directly to
the cAdLIF stack at training/inference time.
"""
times, units, n_channels, n_bins, dyned_levels = args
dense = ssc_to_dense(times, units, n_channels=n_channels, n_bins=n_bins)
# log1p preprocessing (matches dyned_encode_dense's behaviour) so that the
# continuous quantised values share the same dynamic range / scale.
log_dense = np.log1p(dense.astype(np.float32))
quantised = dyned_quantise_2d(log_dense, levels=dyned_levels, boost=False) # [n_channels, n_bins]
if hasattr(quantised, "detach"):
quantised = quantised.detach().cpu().numpy()
# Quantised values live in [0, dyned_levels-1]; for levels<=256 they fit
# in uint8 and the cache becomes 4x smaller than float32. The loader casts
# back to float32 at __getitem__ time.
arr = np.asarray(quantised)
if dyned_levels <= 256:
arr = np.clip(np.rint(arr), 0, 255).astype(np.uint8)
else:
arr = arr.astype(np.float32)
return {"quantised": arr}
def ssc_to_dense(times, units, n_channels=SSC_N_CHANNELS, n_bins=SSC_N_BINS):
"""Convert raw SSC spike events to a dense [n_channels, n_bins] histogram.
Args:
times: float16 array of spike times (seconds)
units: uint16 array of channel indices (0-699)
n_channels: number of cochlea channels (700)
n_bins: number of time bins
Returns:
dense: float32 array of shape [n_channels, n_bins]
"""
times = np.asarray(times, dtype=np.float32)
units = np.asarray(units, dtype=np.int64)
if len(times) == 0:
return np.zeros((n_channels, n_bins), dtype=np.float32)
t_min = times.min()
t_max = times.max()
duration = t_max - t_min
if duration <= 0:
return np.zeros((n_channels, n_bins), dtype=np.float32)
bin_edges = np.linspace(t_min, t_max, n_bins + 1)
bin_indices = np.digitize(times, bin_edges) - 1
bin_indices = np.clip(bin_indices, 0, n_bins - 1)
valid = (units >= 0) & (units < n_channels)
flat_idx = units[valid] * n_bins + bin_indices[valid]
counts = np.bincount(flat_idx, minlength=n_channels * n_bins)
return counts[: n_channels * n_bins].reshape(n_channels, n_bins).astype(np.float32)
# =============================================================================
# ON/OFF Event Encoding Helpers - DyNEDc lossless round-trip via differencing
# (ported verbatim from speech_cadlif_dyned_dynedc_snn_cached-8.py / v8)
# =============================================================================
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)
# =============================================================================
# DyNEDc Compression Measurement (cAdLIF spike outputs at inference)
# =============================================================================
def measure_compression_stats(net, test_loader, device, chunk_size=4):
"""Run DyNEDc over the trained net's cAdLIF spike outputs on the test set.
The cAdLIF layers natively emit binary spike trains during inference;
DyNEDc compresses those spike trains losslessly, which is exactly the
thesis claim ("DyNEDc compresses spike trains"). Stats are aggregated
over each layer's `[T, B, hidden]` spike tensor, per sample.
"""
from collections import Counter
net.eval()
layer_keys = ("spk1", "spk2")
per_layer = {k: {"ratios": [], "modes": []} for k in layer_keys}
with torch.no_grad():
for data, _ in test_loader:
data = data.to(device)
layer_data = net.diagnostic_forward(data)
for k in layer_keys:
spk = layer_data[k] # [T, B, hidden]
spk_np = spk.permute(1, 0, 2).cpu().numpy().astype(np.uint8) # [B, T, hidden]
B = spk_np.shape[0]
for i in range(B):
_, info, _ = dynedc_compress_binary(spk_np[i], chunk_size=chunk_size)
per_layer[k]["ratios"].append(info["compression_ratio"])
per_layer[k]["modes"].append(info.get("mode", "unknown"))
summary = {"chunk_size": chunk_size, "tensor": "cadlif_spike_output"}
all_ratios = []
for k in layer_keys:
ratios = per_layer[k]["ratios"]
modes = per_layer[k]["modes"]
all_ratios.extend(ratios)
summary[k] = {
"mean_compression_ratio": float(np.mean(ratios)),
"mean_space_saving_pct": float((1 - np.mean(ratios)) * 100),
"min_ratio": float(np.min(ratios)),
"max_ratio": float(np.max(ratios)),
"n_samples": len(ratios),
"mode_distribution": dict(Counter(modes)),
}
summary["overall"] = {
"mean_compression_ratio": float(np.mean(all_ratios)),
"mean_space_saving_pct": float((1 - np.mean(all_ratios)) * 100),
"min_ratio": float(np.min(all_ratios)),
"max_ratio": float(np.max(all_ratios)),
"n_samples": len(all_ratios),
}
return summary
# =============================================================================
# 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 SSC SNN with DyNEDc
# =============================================================================
class cAdLIFSSCSNN(nn.Module):
"""SNN with cAdLIF neurons and learnable delays for SSC classification.
Architecture:
Input (700 x T continuous DyNED-quantised values)
-> Linear + BN + Delay + cAdLIF + Drop (layer 1)
-> Linear + BN + Delay + cAdLIF + Drop (layer 2)
-> Linear + LIF readout (no spike, softmax accumulation)
"""
def __init__(self, n_channels=SSC_N_CHANNELS, hidden_size=512, num_outputs=35, max_delay=20, dropout=0.25):
super().__init__()
self.hidden_size = hidden_size
self.num_outputs = num_outputs
self.max_delay = max_delay
# NB: DyNEDc is measured at inference over the test set, compressing
# the binary spike trains the cAdLIF layers emit during their forward
# pass (see `measure_compression_stats`). DyNEDc is informational and
# not in the data path because it is lossless and pre-determined by
# the trained cAdLIF outputs.
# Layer 1: input -> hidden
self.fc1 = nn.Linear(n_channels, hidden_size)
self.bn1 = nn.BatchNorm1d(hidden_size, momentum=0.05)
self.cadlif1 = cAdLIFNeuron(hidden_size)
self.delay1 = nn.Parameter(torch.zeros(hidden_size))
self.drop1 = nn.Dropout(dropout)
# Layer 2: hidden -> hidden
self.fc2 = nn.Linear(hidden_size, hidden_size)
self.bn2 = nn.BatchNorm1d(hidden_size, momentum=0.05)
self.cadlif2 = cAdLIFNeuron(hidden_size)
self.delay2 = nn.Parameter(torch.zeros(hidden_size))
self.drop2 = nn.Dropout(dropout)
# Readout: hidden -> classes (infinite threshold, no spike)
self.fc_out = nn.Linear(hidden_size, num_outputs)
self.alpha_out_raw = nn.Parameter(torch.empty(num_outputs).uniform_(ALPHA_MIN, ALPHA_MAX))
# Weight initialization
nn.init.kaiming_uniform_(self.fc1.weight)
nn.init.kaiming_uniform_(self.fc2.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_channels, n_time_bins]
Returns:
output: [batch, num_outputs] - accumulated softmax votes
"""
B = x.size(0)
T = x.size(2)
device = x.device
# ---- Layer 1: pre-compute linear + BN + delays over full sequence ----
x_seq = x.permute(2, 0, 1) # [T, B, channels]
h1 = self.fc1(x_seq) # [T, B, hidden]
h1 = self.bn1(h1.reshape(-1, self.hidden_size)).reshape(T, B, self.hidden_size)
h1 = apply_delays(h1, self.delay1.clamp(0, self.max_delay), self.max_delay)
# Run cAdLIF layer 1
s1_list = []
u1 = torch.zeros(B, self.hidden_size, device=device)
w1 = torch.zeros(B, self.hidden_size, device=device)
s1 = torch.zeros(B, self.hidden_size, device=device)
for t in range(T):
s1, u1, w1 = self.cadlif1(h1[t], u1, w1, s1)
s1_list.append(s1)
s1_seq = torch.stack(s1_list) # [T, B, hidden]
s1_seq = self.drop1(s1_seq)
# ---- Layer 2: pre-compute linear + BN + delays ----
h2 = self.fc2(s1_seq) # [T, B, hidden]
h2 = self.bn2(h2.reshape(-1, self.hidden_size)).reshape(T, B, self.hidden_size)
h2 = apply_delays(h2, self.delay2.clamp(0, self.max_delay), self.max_delay)
# Run cAdLIF layer 2
s2_list = []
u2 = torch.zeros(B, self.hidden_size, device=device)
w2 = torch.zeros(B, self.hidden_size, device=device)
s2 = torch.zeros(B, self.hidden_size, device=device)
for t in range(T):
s2, u2, w2 = self.cadlif2(h2[t], u2, w2, s2)
s2_list.append(s2)
s2_seq = torch.stack(s2_list) # [T, B, hidden]
s2_seq = self.drop2(s2_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(s2_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 data for monitoring."""
B = x.size(0)
T = x.size(2)
device = x.device
# NB: model receives continuous DyNED-quantised values from the cache.
# DyNEDc is measured at inference on the cAdLIF spike outputs below
# (see `measure_compression_stats`).
x_seq = x.permute(2, 0, 1)
h1 = self.fc1(x_seq)
h1 = self.bn1(h1.reshape(-1, self.hidden_size)).reshape(T, B, self.hidden_size)
h1 = apply_delays(h1, self.delay1.clamp(0, self.max_delay), self.max_delay)
layer_data = {"spk1": [], "spk2": [], "mem_out": [], "mem1": [], "mem2": []}
u1 = w1 = s1 = torch.zeros(B, self.hidden_size, device=device)
for t in range(T):
s1, u1, w1 = self.cadlif1(h1[t], u1, w1, s1)
layer_data["spk1"].append(s1)
layer_data["mem1"].append(u1)
s1_seq = torch.stack(layer_data["spk1"])
h2 = self.fc2(s1_seq)
h2 = self.bn2(h2.reshape(-1, self.hidden_size)).reshape(T, B, self.hidden_size)
h2 = apply_delays(h2, self.delay2.clamp(0, self.max_delay), self.max_delay)
u2 = w2 = s2 = torch.zeros(B, self.hidden_size, device=device)
alpha_out = self._get_alpha_out()
u_out = torch.zeros(B, self.num_outputs, device=device)
for t in range(T):
s2, u2, w2 = self.cadlif2(h2[t], u2, w2, s2)
layer_data["spk2"].append(s2)
layer_data["mem2"].append(u2)
cur = self.fc_out(s2)
u_out = alpha_out * u_out + (1.0 - alpha_out) * cur
layer_data["mem_out"].append(u_out)
for key in layer_data:
layer_data[key] = torch.stack(layer_data[key])
return layer_data
# =============================================================================
# Dataset
# =============================================================================
class DyNEDSSCDataset(torch.utils.data.Dataset):
"""SSC dataset with continuous DyNED-quantised values (Option B cache).
Pipeline (storage):
h5 spike events -> dense histogram [700, n_bins] -> log1p -> DyNED quantise
-> continuous float32 [700, n_bins] -> stored in .pt cache
Pipeline (runtime, in __getitem__):
cached float32 [700, n_bins] -> augmentations -> cAdLIF SSC SNN
Cache format (single .pt file per split):
{
"data": float32 tensor [N, n_channels, n_bins]
"labels": int64 tensor # class indices
}
"""
def __init__(
self,
split="train",
n_bins=SSC_N_BINS,
dyned_levels=256,
ssc_dir=None,
cache_dir=None,
h5_data=None,
dynedc_chunk_size=4,
gen_workers=None,
):
"""
Args:
split: "train", "valid", or "test"
n_bins: number of time bins for dense histogram
dyned_levels: number of sigma-delta quantisation levels
ssc_dir: path to directory containing ssc_{train,valid,test}.h5
cache_dir: path to directory for caching encoded spike trains
dynedc_chunk_size: kept for API compatibility / Optuna param tracking;
does not affect the cached continuous data path.
"""
if not HAS_H5PY:
raise ImportError("h5py is required to load SSC data. Install with: pip install h5py")
self.split = split
self.n_bins = n_bins
self.dyned_levels = dyned_levels
self.is_training = split == "train"
self.dynedc_chunk_size = dynedc_chunk_size
self.gen_workers = gen_workers
self.labels = SSC_LABELS
self.label_to_idx = {label: idx for idx, label in enumerate(self.labels)}
if ssc_dir is None:
ssc_dir = _DEFAULT_SSC_DIR
self.ssc_dir = ssc_dir
if cache_dir is None:
cache_dir = os.path.join("..", "assets")
# Option B cache: continuous DyNED-quantised values (no ON/OFF, no DyNEDc
# at cache time). Filename suffix bumped (`_continuous`) so old Option-A
# caches are not loaded by mistake.
cache_filename = f"dyned_continuous_cache_{split}_bins{n_bins}_lvl{dyned_levels}.pt"
cache_path = Path(cache_dir) / "ssc_cadlif_dyned_dynedc_snn_v2" / cache_filename
self._cache_path = cache_path
if cache_path.exists():
print(f"Loading continuous DyNED cache from {cache_path}...")
try:
cache = torch.load(cache_path, weights_only=False)
self._data = cache["data"]
self.label_indices = cache["labels"]
size_mb = cache_path.stat().st_size / (1024 * 1024)
print(
f"Loaded {len(self._data)} continuous samples "
f"(file: {size_mb:.1f} MB on disk; shape={tuple(self._data.shape)})"
)
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 continuous DyNED values for SSC [{split}]...")
self._build_cache(cache_path, h5_data=h5_data)
def _build_cache(self, cache_path, h5_data=None):
"""Stream-build the cache to bound memory.
- Reads h5 spike events one sample at a time (does not slurp the whole
dataset into Python lists).
- Pre-allocates a single uint8 numpy array of shape [N, n_channels, n_bins]
and writes per-sample results into slices (no list + stack pattern).
- Workers in `multiprocessing.Pool.imap` process samples in chunks; the
result chunks are written into the array and immediately released.
Storage as uint8 keeps memory <= N * 700 * 50 bytes ~ 3.3 GB for the
full SSC train split, vs >= 13 GB if held as float32 in a Python list
plus the additional copy from `np.stack` and `.float()`.
"""
cache_path.parent.mkdir(parents=True, exist_ok=True)
if h5_data is not None:
all_times, all_units, all_labels = h5_data
n_samples = len(all_labels)
print(f" SSC [{self.split}]: {n_samples} samples (pre-read)")
sample_iter = ((all_times[i], all_units[i]) for i in range(n_samples))
label_iter = (all_labels[i] for i in range(n_samples))
else:
h5_path = os.path.join(self.ssc_dir, f"ssc_{self.split}.h5")
if not os.path.exists(h5_path):
raise FileNotFoundError(
f"SSC h5 file not found: {h5_path}\n"
f"Expected files: ssc_train.h5, ssc_valid.h5, ssc_test.h5 in {self.ssc_dir}\n"
"Download via: tonic.datasets.SSC or from https://zenodo.org/record/7426142"
)
h5_handle = h5py.File(h5_path, "r")
n_samples = len(h5_handle["labels"])
print(f" SSC [{self.split}]: {n_samples} samples, streaming h5...")
sample_iter = ((h5_handle["spikes/times"][i], h5_handle["spikes/units"][i]) for i in range(n_samples))
label_iter = (int(h5_handle["labels"][i]) for i in range(n_samples))
# Pre-allocate the cache tensor at uint8 (no float32 list, no np.stack copy).
dtype = np.uint8 if self.dyned_levels <= 256 else np.float32
data_arr = np.empty((n_samples, SSC_N_CHANNELS, self.n_bins), dtype=dtype)
labels_arr = np.empty(n_samples, dtype=np.int64)
n_workers = self.gen_workers if self.gen_workers is not None else max(1, os.cpu_count() - 1)
print(
f" [{self.split}] Encoding {n_samples} samples (continuous DyNED) "
f"using {n_workers} workers, dtype={dtype.__name__}..."
)
def _arg_gen():
for (times_i, units_i), label_i in zip(sample_iter, label_iter):
labels_arr[_arg_gen.idx] = label_i
_arg_gen.idx += 1
yield (times_i, units_i, SSC_N_CHANNELS, self.n_bins, self.dyned_levels)
_arg_gen.idx = 0
import multiprocessing
try:
with multiprocessing.Pool(n_workers) as pool:
for i, result in enumerate(pool.imap(_encode_ssc_sample, _arg_gen(), chunksize=64)):
data_arr[i] = result["quantised"]
if (i + 1) % 5000 == 0 or (i + 1) == n_samples:
pct = (i + 1) / n_samples * 100
print(f" [{self.split}] Encoded {i + 1}/{n_samples} ({pct:.1f}%)")
finally:
if h5_data is None:
h5_handle.close()
self._data = torch.from_numpy(data_arr) # uint8 (or float32) [N, n_channels, n_bins]
self.label_indices = torch.from_numpy(labels_arr) # int64
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)} continuous DyNED samples to {cache_path} "
f"({size_mb:.1f} MB on disk; shape={tuple(self._data.shape)}, dtype={dtype.__name__})"
)
def __len__(self):
return len(self._data)
def __getitem__(self, n):
# Cache is uint8 (or float32 if dyned_levels > 256); cast to float32
# at load so the rest of the pipeline can stay in float.
sample = self._data[n].to(torch.float32) # [n_channels, n_bins]
label_idx = self.label_indices[n]
if self.is_training:
# Time shift augmentation
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)
# Frequency (channel) masking
if torch.rand(1).item() < 0.3:
n_ch = sample.shape[0]
mask_width = torch.randint(1, n_ch // 10 + 1, (1,)).item()
mask_start = torch.randint(0, n_ch - mask_width, (1,)).item()
sample = sample.clone()
sample[mask_start : mask_start + mask_width, :] = 0
# Time masking
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
return sample, label_idx
# =============================================================================
# Data Setup
# =============================================================================
def _stratified_subset_loader(loader, max_samples, num_workers, shuffle):
"""Build a new DataLoader over a stratified-by-class subset of `loader.dataset`.
Selects up to `max_samples // n_classes` samples per class (mirrors the v8
speech-commands per-class quota subset). Uses the cached `label_indices`
from the dataset to avoid scanning samples.
"""
from torch.utils.data import Subset
dataset = loader.dataset
labels = np.asarray(dataset.label_indices)
n_classes = len(SSC_LABELS)
per_class_quota = max(1, max_samples // n_classes)
indices = []
for cls in range(n_classes):
cls_idx = np.where(labels == cls)[0]
if len(cls_idx) > per_class_quota:
cls_idx = cls_idx[:per_class_quota]
indices.extend(cls_idx.tolist())
if len(indices) >= max_samples:
break
indices = sorted(set(indices))[:max_samples]
print(
f" Stratified subset: selected {len(indices)} samples "
f"across {len(set(int(labels[i]) for i in indices))} classes"
)
subset = Subset(dataset, indices)
loader_kwargs = dict(pin_memory=True)
if num_workers > 0:
loader_kwargs["persistent_workers"] = False
loader_kwargs["prefetch_factor"] = 4
return DataLoader(
subset,
batch_size=loader.batch_size,
shuffle=shuffle,
num_workers=num_workers,
**loader_kwargs,
)
def setup_dataloaders(
batch_size=512,
num_workers=4,
n_bins=SSC_N_BINS,
dyned_levels=256,
ssc_dir=None,
cache_dir=None,
dynedc_chunk_size=4,
pin_memory=True,
multiprocessing_context=None,
):
train_dataset = DyNEDSSCDataset(
split="train",
n_bins=n_bins,
dyned_levels=dyned_levels,
ssc_dir=ssc_dir,
cache_dir=cache_dir,
dynedc_chunk_size=dynedc_chunk_size,
)
val_dataset = DyNEDSSCDataset(
split="valid",
n_bins=n_bins,
dyned_levels=dyned_levels,
ssc_dir=ssc_dir,
cache_dir=cache_dir,
dynedc_chunk_size=dynedc_chunk_size,
)
test_dataset = DyNEDSSCDataset(
split="test",
n_bins=n_bins,
dyned_levels=dyned_levels,
ssc_dir=ssc_dir,
cache_dir=cache_dir,
dynedc_chunk_size=dynedc_chunk_size,
)
loader_kwargs = dict(pin_memory=pin_memory)
if num_workers > 0:
loader_kwargs["persistent_workers"] = False
loader_kwargs["prefetch_factor"] = 4
if multiprocessing_context is not None:
loader_kwargs["multiprocessing_context"] = multiprocessing_context
train_loader = DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
**loader_kwargs,
)
val_loader = DataLoader(
val_dataset,
batch_size=batch_size,
shuffle=False,
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, val_loader, test_loader
# =============================================================================
# Training Loop
# =============================================================================
def train_network(
net, train_loader, val_loader, num_epochs=250, device="cuda", lr=1e-3, lr_delay=0.1, weight_decay=1e-4, trial=None
):
# Separate parameter groups:
# - delays get a higher learning rate, no weight decay
# - cAdLIF cell parameters (alpha_raw, beta_raw, a_raw, b_raw) and the
# readout alpha_out_raw are bounded and should not be weight-decayed
# toward zero. Decaying a_raw / b_raw collapses cAdLIF to plain LIF;
# decaying alpha/beta drags time constants out of their valid range.
delay_params = [net.delay1, net.delay2]
neuron_params = [
net.cadlif1.alpha_raw,
net.cadlif1.beta_raw,
net.cadlif1.a_raw,
net.cadlif1.b_raw,
net.cadlif2.alpha_raw,
net.cadlif2.beta_raw,
net.cadlif2.a_raw,
net.cadlif2.b_raw,
net.alpha_out_raw,
]
no_decay_ids = {id(p) for p in delay_params + neuron_params}
weight_params = [p for p in net.parameters() if id(p) not in no_decay_ids]
optimizer = torch.optim.Adam(
[
{"params": weight_params, "lr": lr, "weight_decay": weight_decay},
{"params": neuron_params, "lr": lr, "weight_decay": 0.0},
{"params": delay_params, "lr": lr_delay, "weight_decay": 0.0},
]
)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer,
mode="max",
factor=0.7,
patience=5,
min_lr=1e-6,
)
scaler = torch.amp.GradScaler("cuda") if device == "cuda" else None
best_acc = 0
metrics = {
"epoch": [],
"train_loss": [],
"val_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)
optimizer.zero_grad()
if scaler is not None:
with torch.amp.autocast("cuda"):
output = net(data)
loss = F.cross_entropy(output, targets)
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)
loss = F.cross_entropy(output, targets)
loss.backward()
torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)
optimizer.step()
batch_loss = loss.item()
running_loss += batch_loss
epoch_loss += batch_loss
num_batches += 1
if i % 50 == 49:
avg_loss = running_loss / 50
print(f"Epoch {epoch + 1}, Batch {i + 1}: Loss = {avg_loss:.4f}", end="")
if device == "cuda":
print(f" | GPU: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
else:
print()
running_loss = 0.0
epoch_time = time.time() - epoch_start
avg_epoch_loss = epoch_loss / num_batches
current_lr = optimizer.param_groups[0]["lr"]
# Firing rates diagnostic (use val_loader for monitoring)
with torch.no_grad():
net.eval()
diag_data = next(iter(val_loader))[0][:32].to(device)
layer_data = net.diagnostic_forward(diag_data)
firing_rates = {
"layer1": layer_data["spk1"].mean().item(),
"layer2": layer_data["spk2"].mean().item(),
}
eval_result = evaluate(net, val_loader, device)
val_acc = eval_result["accuracy"]
# Step scheduler with validation accuracy
scheduler.step(val_acc)
metrics["epoch"].append(epoch + 1)
metrics["train_loss"].append(avg_epoch_loss)
metrics["val_accuracy"].append(val_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
alpha1 = net.cadlif1.alpha_raw.clamp(ALPHA_MIN, ALPHA_MAX).mean().item()
alpha2 = net.cadlif2.alpha_raw.clamp(ALPHA_MIN, ALPHA_MAX).mean().item()
d1_mean = net.delay1.clamp(0, net.max_delay).mean().item()
d2_mean = net.delay2.clamp(0, net.max_delay).mean().item()
print(
f"Epoch {epoch + 1}: Val Acc = {val_acc:.4f} | Loss = {avg_epoch_loss:.4f} | "
f"LR = {current_lr:.6f} | alpha = [{alpha1:.3f}, {alpha2:.3f}] | "
f"Delay = [{d1_mean:.1f}, {d2_mean:.1f}] | FR = [{firing_rates['layer1']:.3f}, "
f"{firing_rates['layer2']:.3f}] | {epoch_time:.1f}s"
)
if device == "cuda":
torch.cuda.empty_cache()
if val_acc > best_acc:
best_acc = val_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": val_acc,
},
os.path.join(OUTPUT_DIR, "best_cadlif_ssc_dynedc_model.pth"),
)
if trial is not None:
import optuna
trial.report(val_acc, epoch)
if trial.should_prune():
raise optuna.TrialPruned()
# Restore best model weights before returning
if best_model_state is not None:
net.load_state_dict(best_model_state)
return metrics
# =============================================================================
# Evaluation
# =============================================================================
def evaluate(net, loader, device, collect_representations=False):
net.eval()
correct = 0
total = 0
all_predictions = []
all_targets = []
all_representations = []
num_classes = len(SSC_LABELS)
per_class_correct = np.zeros(num_classes)
per_class_total = np.zeros(num_classes)
with torch.no_grad():
for data, targets in 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": {SSC_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 measured separately by
# `measure_compression_stats`, which compresses the cAdLIF spike outputs
# of the trained net.
return result
# =============================================================================
# Visualizations
# =============================================================================
def analyze_training_metrics(metrics):
losses = np.array(metrics["train_loss"])
accuracies = np.array(metrics["val_accuracy"])
epochs = np.array(metrics["epoch"])
# Save CSV
header_parts = ["epoch", "train_loss", "val_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 SSC_LABELS)
header = ",".join(header_parts)
rows = []
for i in range(len(metrics["epoch"])):
row = [
metrics["epoch"][i],
metrics["train_loss"][i],
metrics["val_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 SSC_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} - validation 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="Validation Accuracy", alpha=0.8)
ax2.tick_params(axis="y", labelcolor="tab:red")
plt.title("cAdLIF+DyNEDc SSC 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))
rate_arrays = {}
for key in keys:
rates = [fr[key] for fr in fr_history]
plt.plot(epochs, rates, label=key, linewidth=1.5)
rate_arrays[f"rates_{key}"] = np.array(rates)
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),
**rate_arrays,
)
# 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(SSC_LABELS)
fig, ax = plt.subplots(figsize=(14, 12))
cm_normalized = cm.astype(float) / (cm.sum(axis=1, keepdims=True) + 1e-8)
im = ax.imshow(cm_normalized, interpolation="nearest", cmap="Blues")
ax.figure.colorbar(im, ax=ax)
ax.set(
xticks=range(num_classes),
yticks=range(num_classes),
xticklabels=SSC_LABELS,
yticklabels=SSC_LABELS,
ylabel="True Label",
xlabel="Predicted Label",
title="Confusion Matrix (Normalized)",
)
plt.setp(ax.get_xticklabels(), rotation=90, ha="right", rotation_mode="anchor", fontsize=7)
plt.setp(ax.get_yticklabels(), fontsize=7)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "confusion_matrix.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR,
"confusion_matrix",
cm=cm,
cm_normalized=cm_normalized,
labels=np.array(SSC_LABELS),
)
np.savetxt(
os.path.join(OUTPUT_DIR, "confusion_matrix.csv"),
cm,
delimiter=",",
fmt="%d",
header=",".join(SSC_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(SSC_LABELS)))
cbar.set_ticklabels(SSC_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(SSC_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))
bars = plt.bar(range(len(labels_sorted)), [a * 100 for a in accs_sorted], color="steelblue")
plt.xticks(range(len(labels_sorted)), labels_sorted, rotation=90, 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",
accuracies_sorted=np.array(accs_sorted),
labels_sorted=np.array(labels_sorted),
labels=np.array(labels),
accuracies=np.array(accs),
)
print("Saved per-class accuracy plot")
def collect_diagnostic_batches(net, loader, device, max_samples=4000):
"""Run diagnostic_forward across multiple 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 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, last hidden
raster, last hidden firing-rate distribution, output membrane potentials."""
spk_last = layer_data["spk2"] # [T, B, hidden] - last hidden layer
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="viridis")
axes[0, 0].set_title("Input DyNED-quantised values (Sample 0)")
axes[0, 0].set_xlabel("Time Bin")
# Input is the continuous DyNED-quantised array [n_channels, T].
axes[0, 0].set_ylabel("Cochlea Channel")
axes[0, 1].imshow(spk_last[:, 0].cpu().numpy(), aspect="auto", cmap="binary")
axes[0, 1].set_title("Hidden Layer 2 Spike Raster (Sample 0)")
axes[0, 1].set_xlabel("Neuron Index")
axes[0, 1].set_ylabel("Time Step")
rates = spk_last.mean(dim=0).cpu().numpy().flatten()
axes[1, 0].hist(rates, bins=50, color="steelblue")
axes[1, 0].set_title("Hidden Layer 2 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_last_sample0=spk_last[:, 0],
firing_rates=rates,
mem_out_sample0=mem_out[:, 0],
)
print("Saved network activity")
def visualize_layer_spike_rasters(layer_data):
"""Three-panel raster: spk1, spk2, mem_out (heatmap) for sample 0."""
fig, axes = plt.subplots(1, 3, figsize=(18, 6))
panels = [
("spk1", "Layer 1 (Hidden)"),
("spk2", "Layer 2 (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"],
mem_out=panel_arrs["mem_out"],
)
print("Saved layer spike rasters")
def visualize_membrane_distributions(layer_data):
"""Membrane potential histograms for each hidden cAdLIF layer plus the
output-layer membrane distribution."""
fig, axes = plt.subplots(1, 3, figsize=(18, 5))
panels = [
("mem1", "Layer 1 Membrane Potentials"),
("mem2", "Layer 2 Membrane Potentials"),
("mem_out", "Output Membrane Potentials"),
]
panel_data = {}
for ax, (key, title) in zip(axes.flat, panels):
d = layer_data[key].cpu().numpy()
vals = d.flatten()
panel_data[key] = vals
color = "darkorange" if key == "mem_out" else "steelblue"
ax.hist(vals, bins=100, density=True, color=color, 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",
mem1_values=panel_data["mem1"],
mem2_values=panel_data["mem2"],
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()
rows, cols = 5, 7
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(SSC_LABELS), T_dim, C_dim), np.nan, dtype=np.float32)
for cls_idx, ax in enumerate(axes.flat):
if cls_idx >= len(SSC_LABELS):
ax.axis("off")
continue
mask = targets == cls_idx
if mask.sum() == 0:
ax.set_title(SSC_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(SSC_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(SSC_LABELS),
grid=np.array([rows, cols]),
)
print("Saved per-class spikes")
def visualize_weight_distributions(net):
"""Trained weight histograms for the learned linear layers."""
panels = [
("fc1", net.fc1.weight),
("fc2", net.fc2.weight),
("fc_out", net.fc_out.weight),
]
n = len(panels)
fig, axes = plt.subplots(1, n, figsize=(5 * n, 5))
if n == 1:
axes = [axes]
weight_payload = {}
name_list = []
shape_list = []
for ax, (name, w) in zip(axes, panels):
vals = w.detach().cpu().numpy().flatten()
weight_payload[f"weight__{name}"] = vals
name_list.append(name)
shape_list.append(list(w.shape))
ax.hist(vals, bins=100, density=True, color="steelblue", alpha=0.8)
ax.set_title(f"{name} {tuple(w.shape)}")
ax.set_xlabel("Weight Value")
ax.set_ylabel("Density")
ax.text(
0.02,
0.98,
f"Mean: {vals.mean():.3f}\nStd: {vals.std():.3f}",
transform=ax.transAxes,
va="top",
fontsize=9,
bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.8),
)
plt.suptitle("Trained Weight Distributions", fontsize=14)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "weight_distributions.png"), dpi=150)
plt.close()
dump_plot_data(
OUTPUT_DIR, "weight_distributions", names=np.array(name_list), shapes=np.array(shape_list), **weight_payload
)
print("Saved weight distributions")
# =============================================================================
# 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_channels = SSC_N_CHANNELS # 700 cochlea channels
n_bins = SSC_N_BINS # 50 time bins (~20ms per bin)
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", 1e-4)
hidden_size = params["hidden_size"]
max_delay = params.get("max_delay", 20)
dropout = params.get("dropout", 0.25)
batch_size = params.get("batch_size", 512)
dyned_levels = params.get("dyned_levels", 32)
dynedc_chunk_size = params.get("dynedc_chunk_size", 4)
n_bins = params.get("n_bins", SSC_N_BINS)
else:
lr = 1e-3
lr_delay = 0.1
weight_decay = 1e-4
hidden_size = 512
max_delay = 20
dropout = 0.25
batch_size = 512
dyned_levels = 32
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/val/test datasets to ~{max_samples} stratified samples")
print("=" * 80)
print("DyNED -> cAdLIF SNN for Spiking Speech Commands (SSC v2) + DyNEDc on cAdLIF spike outputs at inference")
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: SSC cochlea spikes -> dense histogram ({n_channels}ch x {n_bins}bins)"
f" -> DyNED (lvl={dyned_levels}) continuous values cached on disk"
)
print(f"Runtime: cached continuous DyNED values [{n_channels}, {n_bins}] -> cAdLIF SNN")
print(
f"Inference: DyNEDc (chunk={dynedc_chunk_size}) compresses the binary spike trains "
f"the cAdLIF layers emit during inference (spk1, spk2)"
)
print(f"Input: {n_channels} continuous DyNED-quantised features x {n_bins} time bins")
print(f"Architecture: {n_channels} -> {hidden_size} -> {hidden_size} -> {num_outputs}")
print(f"Neurons: cAdLIF (learnable alpha, beta, a, b) + delays (max={max_delay})")
print(f"Batch: {batch_size} | Hidden: {hidden_size} | Dropout: {dropout}")
print(f"LR: {lr} (weights) / {lr_delay} (delays) | Weight decay: {weight_decay}")
print("=" * 80)
net = cAdLIFSSCSNN(
n_channels=SSC_N_CHANNELS,
hidden_size=hidden_size,
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, val_loader, test_loader = setup_dataloaders(
batch_size=batch_size,
num_workers=num_workers,
n_bins=n_bins,
dyned_levels=dyned_levels,
dynedc_chunk_size=dynedc_chunk_size,
)
if max_samples is not None:
# Stratified subset by class label (mirrors v8's per-class quota approach).
train_loader = _stratified_subset_loader(train_loader, max_samples, num_workers, shuffle=True)
val_loader = _stratified_subset_loader(val_loader, max_samples, num_workers, shuffle=False)
test_loader = _stratified_subset_loader(test_loader, max_samples, num_workers, shuffle=False)
try:
metrics = train_network(
net=net,
train_loader=train_loader,
val_loader=val_loader,
num_epochs=num_epochs,
device=device_str,
lr=lr,
lr_delay=lr_delay,
weight_decay=weight_decay,
)
print("\n" + "=" * 80)
print("INFERENCE ANALYSIS")
print("=" * 80)
analyze_training_metrics(metrics)
# Final evaluation on test set (best model restored in train_network)
print("\nEvaluating best model on test set...")
test_eval = evaluate(net, test_loader, device, collect_representations=True)
print(f"Test accuracy: {test_eval['accuracy']:.4f}")
# DyNEDc compression on the cAdLIF spike outputs at inference: forward
# the trained net over the test set, compress each sample's spk1 / spk2
# with DyNEDc. DyNEDc is lossless, so this number reflects the on-wire /
# on-disk savings without affecting model accuracy.
print(f"\nMeasuring DyNEDc compression on cAdLIF spike outputs (chunk_size={dynedc_chunk_size})...")
stats = measure_compression_stats(
net,
test_loader,
device,
chunk_size=dynedc_chunk_size,
)
print("DyNEDc compression statistics (test set, cAdLIF spike outputs):")
for layer_key in ("spk1", "spk2", "overall"):
ls = stats[layer_key]
print(
f" {layer_key:<8}: ratio = {ls['mean_compression_ratio']:.4f} "
f"saving = {ls['mean_space_saving_pct']:.2f}% "
f"range = [{ls['min_ratio']:.4f}, {ls['max_ratio']:.4f}] "
f"n = {ls['n_samples']}"
)
import json as _json
with open(os.path.join(OUTPUT_DIR, "dynedc_compression_stats.json"), "w") as _fh:
_json.dump(stats, _fh, indent=2)
visualize_confusion_matrix_plot(test_eval["confusion_matrix"])
if "representations" in test_eval:
visualize_tsne(test_eval["representations"], test_eval["targets"])
visualize_per_class_accuracy(test_eval["per_class_accuracy"])
# Network dynamics figures from a 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 (test set):")
for cls, acc in sorted(test_eval["per_class_accuracy"].items(), key=lambda x: x[1], reverse=True):
print(f" {cls:>12s}: {acc:.4f}")
# Print learned neuron parameters
print("\nLearned neuron parameters:")
for i, cadlif in enumerate([net.cadlif1, net.cadlif2], 1):
alpha, beta, a, b = cadlif._constrain()
print(f" Layer {i}: 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:")
print(
f" Layer 1: mean={net.delay1.clamp(0, net.max_delay).mean():.2f}, "
f"std={net.delay1.clamp(0, net.max_delay).std():.2f}, "
f"range=[{net.delay1.clamp(0, net.max_delay).min():.1f}, "
f"{net.delay1.clamp(0, net.max_delay).max():.1f}]"
)
print(
f" Layer 2: mean={net.delay2.clamp(0, net.max_delay).mean():.2f}, "
f"std={net.delay2.clamp(0, net.max_delay).std():.2f}, "
f"range=[{net.delay2.clamp(0, net.max_delay).min():.1f}, "
f"{net.delay2.clamp(0, net.max_delay).max():.1f}]"
)
torch.save(
{
"model_state_dict": net.state_dict(),
"dyned_levels": dyned_levels,
"n_channels": n_channels,
"n_bins": n_bins,
"dynedc_chunk_size": dynedc_chunk_size,
"metrics": {
"train_losses": metrics["train_loss"],
"val_accuracies": metrics["val_accuracy"],
"test_accuracy": test_eval["accuracy"],
},
},
os.path.join(OUTPUT_DIR, "final_model.pth"),
)
print(
f"\nDone! Best val accuracy: {max(metrics['val_accuracy']):.4f} | "
f"Test accuracy: {test_eval['accuracy']:.4f}"
)
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_channels = SSC_N_CHANNELS
num_outputs = 35
lr = trial.suggest_float("lr", 1e-4, 2e-3, log=True)
lr_delay = trial.suggest_float("lr_delay", 1e-2, 1.0, log=True)
weight_decay = trial.suggest_float("weight_decay", 1e-5, 1e-2, log=True)
hidden_size = trial.suggest_categorical("hidden_size", [256, 512])
max_delay = trial.suggest_categorical("max_delay", [10, 15, 20, 25])
dropout = trial.suggest_float("dropout", 0.1, 0.5)
batch_size = trial.suggest_categorical("batch_size", [256, 512, 1024])
dyned_levels = trial.suggest_categorical("dyned_levels", [16, 32, 64, 128, 256])
n_bins = trial.suggest_categorical("n_bins", [25, 50, 100])
# 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
# Default to 4 workers to move uint8->float32 cast and DataLoader collation
# off the main thread (which is otherwise GIL-pegged submitting the cAdLIF
# temporal-loop kernels). spawn context is required because fork + an
# already-initialised CUDA context corrupts the workers.
num_workers = 4
if _CLI_WORKERS is not None:
num_workers = _CLI_WORKERS
net = cAdLIFSSCSNN(
n_channels=SSC_N_CHANNELS,
hidden_size=hidden_size,
num_outputs=num_outputs,
max_delay=max_delay,
dropout=dropout,
).to(device)
# pin_memory=False under Optuna: pinned host memory accumulates across
# trials in the same process (empty_cache only frees CUDA-side memory),
# which eventually causes pin_memory() to fail with cudaErrorInvalidValue.
train_loader, val_loader, test_loader = setup_dataloaders(
batch_size=batch_size,
num_workers=num_workers,
n_bins=n_bins,
dyned_levels=dyned_levels,
dynedc_chunk_size=dynedc_chunk_size,
pin_memory=False,
multiprocessing_context="spawn" if num_workers > 0 else None,
)
try:
metrics = train_network(
net=net,
train_loader=train_loader,
val_loader=val_loader,
num_epochs=num_epochs,
device=device_str,
lr=lr,
lr_delay=lr_delay,
weight_decay=weight_decay,
trial=trial,
)
except optuna.TrialPruned:
raise
finally:
# Tear down workers + datasets fully so the next trial starts clean.
for ld in (train_loader, val_loader, test_loader):
it = getattr(ld, "_iterator", None)
if it is not None and hasattr(it, "_shutdown_workers"):
it._shutdown_workers()
del net, train_loader, val_loader, test_loader
import gc
gc.collect()
if device_str == "cuda":
torch.cuda.empty_cache()
torch.cuda.synchronize()
return max(metrics["val_accuracy"])
def optuna_optimize(n_trials=50, study_name="ssc_cadlif_dyned_dynedc_snn", 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 validation 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="SSC cAdLIF DyNED+DyNEDc SNN")
parser.add_argument("--optuna", action="store_true", help="Run Optuna hyperparameter optimisation")
parser.add_argument("--n-trials", type=int, default=50, help="Number of Optuna trials (default: 50)")
parser.add_argument("--study-name", type=str, default="ssc_cadlif_dyned_dynedc_snn", help="Optuna study name")
parser.add_argument("--n-jobs", type=int, default=1, help="Number of parallel Optuna jobs")
parser.add_argument("--json", type=str, default=None, help="Path to JSON file with best hyperparameters")
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("--workers", type=int, default=None, help="DataLoader/cache num_workers (default: auto)")
parser.add_argument(
"--gen-workers",
type=int,
default=None,
help="Worker count for --gen-data multiprocessing.Pool (default: cpu_count - 1)",
)
parser.add_argument(
"--subset-size",
type=int,
default=None,
help="Limit train/val/test loaders to N stratified 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()
if args.cpu:
os.environ["CUDA_VISIBLE_DEVICES"] = ""
_CLI_WORKERS = args.workers
_CLI_SUBSET_SIZE = args.subset_size
_CLI_QUICK_EPOCHS = args.quick_epochs
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 [16, 32, 64, 128, 256]:
for n_bins in [25, 50, 100]:
for split in ["train", "valid", "test"]:
# Option B caches store continuous DyNED-quantised values;
# they do not depend on dynedc_chunk_size (chunk_size only
# affects the inference-time spike-output compression).
prefix = f"lvl={levels},bins={n_bins},{split}"
combos.append((prefix, dict(split=split, n_bins=n_bins, dyned_levels=levels)))
print(f"Pre-generating {len(combos)} SSC cache files...")
# Verify h5 files exist before starting (they will be streamed inside
# _build_cache, never slurped). Slurping all 3 splits up-front pushes
# past 60 GB RAM on large host configs; streaming keeps the build to
# the data_arr footprint (<= ~5 GB for train at bins=100, uint8) plus
# the worker pool overhead.
for split in ["train", "valid", "test"]:
h5_path = os.path.join(_DEFAULT_SSC_DIR, f"ssc_{split}.h5")
if not os.path.exists(h5_path):
print(f"ERROR: SSC h5 file not found: {h5_path}")
print(f"Download via: tonic.datasets.SSC or from https://zenodo.org/record/7426142")
sys.exit(1)
import gc
for prefix, kwargs in combos:
kwargs["gen_workers"] = args.gen_workers
_build_with_prefix(prefix, DyNEDSSCDataset, kwargs)
gc.collect()
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)