ssc_cadlif_dyned_snn_v3.py
SSC DyNED SNN - 1,474 lines.
View on GitHub (speech-neuro/ssc_cadlif_dyned_snn_v3.py).
Source
Section titled “Source”"""
SSC (Spiking Speech Commands) SNN with DyNED Encoding - cAdLIF + Learnable Delays (v3)
Pipeline: SSC cochlea spikes -> dense spike histogram (700ch x bins) -> DyNED
(sigma-delta quantisation) -> cAdLIF SNN
Architecture based on:
- Deckers et al. (2024) "Co-learning synaptic delays, weights and adaptation"
- Bittar & Garner (2022) "A surrogate gradient spiking baseline for speech command recognition"
v3 changes vs v2:
- Wider ALPHA_MIN = exp(-1/2) ~ 0.607 (was 0.819) - v2 alphas all saturated at floor
- Label smoothing 0.1 - reduces overfitting (v2 train loss dropped but val acc flat)
- Stronger augmentation - masking prob 0.5 (was 0.3), wider masks
Carries forward from v2:
- 3 hidden cAdLIF layers, CosineAnnealingLR, n_bins=100
Run: uv run python speech-neuro/ssc_cadlif_dyned_snn_v3.py
Optuna: uv run python speech-neuro/ssc_cadlif_dyned_snn_v3.py --optuna --n-trials 50
"""
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, dyned_encode_dense
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_output_v3")
os.makedirs(OUTPUT_DIR, exist_ok=True)
_CLI_WORKERS = None # Set by --workers flag
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 ~10ms per bin over ~1s)
SSC_N_BINS = 100
# =============================================================================
# SSC Dense Histogram + DyNED Encoding
# =============================================================================
def _encode_ssc_sample(args):
"""Worker function for multiprocessing: encode a single SSC sample."""
times, units, n_channels, n_bins, dyned_levels = args
dense = ssc_to_dense(times, units, n_channels=n_channels, n_bins=n_bins)
spike_train = dyned_encode_dense(dense, dyned_levels=dyned_levels)
return spike_train.numpy()
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)
# =============================================================================
# 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 / 2.0) # 0.6065 - fast membrane leak (widened from 0.8187, v2 saturated at floor)
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.6065, 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
# =============================================================================
class cAdLIFSSCSNN(nn.Module):
"""SNN with cAdLIF neurons and learnable delays for SSC classification.
Architecture (3 hidden layers, matching Deckers et al. 2024):
Input (700 x T) -> Linear + BN + Delay + cAdLIF + Drop (layer 1)
-> Linear + BN + Delay + cAdLIF + Drop (layer 2)
-> Linear + BN + Delay + cAdLIF + Drop (layer 3)
-> 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
# 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)
# Layer 3: hidden -> hidden
self.fc3 = nn.Linear(hidden_size, hidden_size)
self.bn3 = nn.BatchNorm1d(hidden_size, momentum=0.05)
self.cadlif3 = cAdLIFNeuron(hidden_size)
self.delay3 = nn.Parameter(torch.zeros(hidden_size))
self.drop3 = 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.fc3.weight)
nn.init.kaiming_uniform_(self.fc_out.weight)
def _get_alpha_out(self):
return self.alpha_out_raw.clamp(ALPHA_MIN, ALPHA_MAX)
def 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)
# ---- Layer 3: pre-compute linear + BN + delays ----
h3 = self.fc3(s2_seq) # [T, B, hidden]
h3 = self.bn3(h3.reshape(-1, self.hidden_size)).reshape(T, B, self.hidden_size)
h3 = apply_delays(h3, self.delay3.clamp(0, self.max_delay), self.max_delay)
# Run cAdLIF layer 3
s3_list = []
u3 = torch.zeros(B, self.hidden_size, device=device)
w3 = torch.zeros(B, self.hidden_size, device=device)
s3 = torch.zeros(B, self.hidden_size, device=device)
for t in range(T):
s3, u3, w3 = self.cadlif3(h3[t], u3, w3, s3)
s3_list.append(s3)
s3_seq = torch.stack(s3_list) # [T, B, hidden]
s3_seq = self.drop3(s3_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(s3_seq[t])
u_out = alpha_out * u_out + (1.0 - alpha_out) * cur
m_out_list.append(u_out)
m_out = torch.stack(m_out_list) # [T, B, 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
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": [], "spk3": [], "mem_out": []}
u1 = w1 = s1 = torch.zeros(B, self.hidden_size, device=device)
for t in range(T):
s1, u1, w1 = self.cadlif1(h1[t], u1, w1, s1)
layer_data["spk1"].append(s1)
s1_seq = torch.stack(layer_data["spk1"])
h2 = self.fc2(s1_seq)
h2 = self.bn2(h2.reshape(-1, self.hidden_size)).reshape(T, B, self.hidden_size)
h2 = apply_delays(h2, self.delay2.clamp(0, self.max_delay), self.max_delay)
u2 = w2 = s2 = torch.zeros(B, self.hidden_size, device=device)
for t in range(T):
s2, u2, w2 = self.cadlif2(h2[t], u2, w2, s2)
layer_data["spk2"].append(s2)
s2_seq = torch.stack(layer_data["spk2"])
h3 = self.fc3(s2_seq)
h3 = self.bn3(h3.reshape(-1, self.hidden_size)).reshape(T, B, self.hidden_size)
h3 = apply_delays(h3, self.delay3.clamp(0, self.max_delay), self.max_delay)
u3 = w3 = s3 = torch.zeros(B, self.hidden_size, device=device)
alpha_out = self._get_alpha_out()
u_out = torch.zeros(B, self.num_outputs, device=device)
for t in range(T):
s3, u3, w3 = self.cadlif3(h3[t], u3, w3, s3)
layer_data["spk3"].append(s3)
cur = self.fc_out(s3)
u_out = alpha_out * u_out + (1.0 - alpha_out) * cur
layer_data["mem_out"].append(u_out)
for key in layer_data:
layer_data[key] = torch.stack(layer_data[key])
return layer_data
# =============================================================================
# Dataset
# =============================================================================
class DyNEDSSCDataset(torch.utils.data.Dataset):
"""SSC dataset with pre-computed DyNED spike trains cached to disk.
Reads SSC h5 files directly to avoid Tonic's float16 timestamp overflow bug.
Pipeline: h5 spike events -> dense histogram [700, n_bins] -> DyNED -> cached.
"""
def __init__(
self,
split="train",
n_bins=SSC_N_BINS,
dyned_levels=256,
ssc_dir=None,
cache_dir=None,
h5_data=None,
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
gen_workers: pool size for cache-build multiprocessing
(default: cpu_count - 1)
"""
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.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")
cache_path = (
Path(cache_dir) / "ssc_cadlif_dyned_snn_v3" / f"ssc_dyned_cache_{split}_bins{n_bins}_lvl{dyned_levels}.pt"
)
if cache_path.exists():
print(f"Loading cached SSC spike trains from {cache_path}...")
try:
cache = torch.load(cache_path, weights_only=True)
self.spike_trains = cache["spike_trains"]
self.label_indices = cache["labels"]
print(f"Loaded {len(self.spike_trains)} samples ({self.spike_trains.shape})")
return
except (RuntimeError, EOFError, zipfile.BadZipFile) as e:
print(f"Corrupted cache file, rebuilding: {e}")
cache_path.unlink(missing_ok=True)
print(f"Cache not found - pre-computing DyNED spike trains 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 (or uses the pre-read
h5_data tuple if provided), pre-allocates a single float32 numpy
array of shape [N, n_channels, n_bins], and writes per-sample
results into slices via multiprocessing.imap. Avoids the
list + torch.stack pattern that previously held two full copies of
the dataset in RAM at the same time.
"""
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))
h5_handle = None
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))
spike_trains_arr = np.empty((n_samples, SSC_N_CHANNELS, self.n_bins), dtype=np.float32)
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 using {n_workers} workers...")
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, spike_train_np in enumerate(pool.imap(_encode_ssc_sample, _arg_gen(), chunksize=64)):
spike_trains_arr[i] = spike_train_np
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} samples ({pct:.1f}%)")
finally:
if h5_handle is not None:
h5_handle.close()
self.spike_trains = torch.from_numpy(spike_trains_arr)
self.label_indices = torch.from_numpy(labels_arr)
torch.save({"spike_trains": self.spike_trains, "labels": self.label_indices}, cache_path)
size_mb = cache_path.stat().st_size / (1024 * 1024)
print(f"Cached {len(self.spike_trains)} spike trains to {cache_path} ({size_mb:.1f} MB)")
def __len__(self):
return len(self.spike_trains)
def __getitem__(self, n):
spike_train = self.spike_trains[n]
label_idx = self.label_indices[n]
if self.is_training:
# Time shift augmentation
max_shift = spike_train.shape[1] // 10
shift = torch.randint(-max_shift, max_shift + 1, (1,)).item()
if shift != 0:
spike_train = torch.roll(spike_train, shifts=shift, dims=1)
# Frequency (channel) masking
if torch.rand(1).item() < 0.5:
n_ch = spike_train.shape[0]
mask_width = torch.randint(1, n_ch // 5 + 1, (1,)).item()
mask_start = torch.randint(0, n_ch - mask_width, (1,)).item()
spike_train = spike_train.clone()
spike_train[mask_start : mask_start + mask_width, :] = 0
# Time masking
if torch.rand(1).item() < 0.5:
n_time = spike_train.shape[1]
mask_width = torch.randint(1, n_time // 5 + 1, (1,)).item()
mask_start = torch.randint(0, n_time - mask_width, (1,)).item()
spike_train = spike_train.clone()
spike_train[:, mask_start : mask_start + mask_width] = 0
return spike_train, label_idx
# =============================================================================
# Data Setup
# =============================================================================
def setup_dataloaders(
batch_size=512,
num_workers=4,
n_bins=SSC_N_BINS,
dyned_levels=256,
ssc_dir=None,
cache_dir=None,
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,
)
val_dataset = DyNEDSSCDataset(
split="valid",
n_bins=n_bins,
dyned_levels=dyned_levels,
ssc_dir=ssc_dir,
cache_dir=cache_dir,
)
test_dataset = DyNEDSSCDataset(
split="test",
n_bins=n_bins,
dyned_levels=dyned_levels,
ssc_dir=ssc_dir,
cache_dir=cache_dir,
)
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 higher learning rate
delay_params = [net.delay1, net.delay2, net.delay3]
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.Adam(
[
{"params": weight_params, "lr": lr, "weight_decay": weight_decay},
{"params": delay_params, "lr": lr_delay, "weight_decay": 0.0},
]
)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer,
T_max=num_epochs,
eta_min=1e-6,
)
scaler = torch.amp.GradScaler("cuda") if device == "cuda" else None
best_acc = 0
metrics = {
"epoch": [],
"train_loss": [],
"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, label_smoothing=0.1)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
else:
output = net(data)
loss = F.cross_entropy(output, targets, label_smoothing=0.1)
loss.backward()
torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)
optimizer.step()
batch_loss = loss.item()
running_loss += batch_loss
epoch_loss += batch_loss
num_batches += 1
if i % 50 == 49:
avg_loss = running_loss / 50
print(f"Epoch {epoch + 1}, Batch {i + 1}: Loss = {avg_loss:.4f}", end="")
if device == "cuda":
print(f" | GPU: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
else:
print()
running_loss = 0.0
epoch_time = time.time() - epoch_start
avg_epoch_loss = epoch_loss / num_batches
current_lr = optimizer.param_groups[0]["lr"]
# Firing rates diagnostic (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(),
"layer3": layer_data["spk3"].mean().item(),
}
eval_result = evaluate(net, val_loader, device)
val_acc = eval_result["accuracy"]
scheduler.step()
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()
alpha3 = net.cadlif3.alpha_raw.clamp(ALPHA_MIN, ALPHA_MAX).mean().item()
d1_mean = net.delay1.clamp(0, net.max_delay).mean().item()
d2_mean = net.delay2.clamp(0, net.max_delay).mean().item()
d3_mean = net.delay3.clamp(0, net.max_delay).mean().item()
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}, {alpha3:.3f}] | "
f"Delay = [{d1_mean:.1f}, {d2_mean:.1f}, {d3_mean:.1f}] | "
f"FR = [{firing_rates['layer1']:.3f}, {firing_rates['layer2']:.3f}, "
f"{firing_rates['layer3']:.3f}] | {epoch_time:.1f}s"
)
if device == "cuda":
torch.cuda.empty_cache()
if 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_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()
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 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()
# Firing rates
if metrics["layer_firing_rates"]:
fr_history = metrics["layer_firing_rates"]
keys = sorted(fr_history[0].keys())
plt.figure(figsize=(12, 6))
for key in keys:
rates = [fr[key] for fr in fr_history]
plt.plot(epochs, rates, label=key, linewidth=1.5)
plt.xlabel("Epoch")
plt.ylabel("Mean Firing Rate")
plt.title("Per-Layer Firing Rates")
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig(os.path.join(OUTPUT_DIR, "firing_rates.png"), dpi=150)
plt.close()
# 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()
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()
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()
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()
print("Saved per-class accuracy plot")
# =============================================================================
# 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 # 100 time bins (~10ms 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)
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
num_epochs = 250
num_workers = min(8, os.cpu_count() - 2) if os.cpu_count() > 2 else 0
print("=" * 80)
print("DyNED + cAdLIF SNN for Spiking Speech Commands (SSC) - v3 (wider alpha, label smoothing)")
print("=" * 80)
print(f"Device: {device_str.upper()}")
if torch.cuda.is_available():
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(
f"Pipeline: SSC cochlea spikes -> dense histogram ({n_channels}ch x {n_bins}bins)"
f" -> DyNED (lvl={dyned_levels}) -> cAdLIF SNN"
)
print(f"Input: {n_channels} channels x {n_bins} time bins")
print(f"Architecture: {n_channels} -> {hidden_size} -> {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=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,
)
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("POST-TRAINING 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}")
visualize_confusion_matrix_plot(test_eval["confusion_matrix"])
if "representations" in test_eval:
visualize_tsne(test_eval["representations"], test_eval["targets"])
visualize_per_class_accuracy(test_eval["per_class_accuracy"])
print("\nFinal Per-Class Accuracy (test set):")
for cls, acc in sorted(test_eval["per_class_accuracy"].items(), key=lambda x: x[1], reverse=True):
print(f" {cls:>12s}: {acc:.4f}")
# Print learned neuron parameters
print("\nLearned neuron parameters:")
for i, cadlif in enumerate([net.cadlif1, net.cadlif2, net.cadlif3], 1):
alpha, beta, a, b = cadlif._constrain()
print(f" Layer {i}: 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, delay in enumerate([net.delay1, net.delay2, net.delay3], 1):
d = delay.clamp(0, net.max_delay)
print(f" Layer {i}: mean={d.mean():.2f}, std={d.std():.2f}, range=[{d.min():.1f}, {d.max():.1f}]")
torch.save(
{
"model_state_dict": net.state_dict(),
"dyned_levels": dyned_levels,
"n_channels": n_channels,
"n_bins": n_bins,
"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"Outputs: {OUTPUT_DIR}")
except KeyboardInterrupt:
print("Interrupted - saving...")
torch.save({"model_state_dict": net.state_dict()}, os.path.join(OUTPUT_DIR, "interrupted_model.pth"))
def optuna_objective(trial):
import optuna
device_str = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device_str)
n_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])
num_epochs = 50
# num_workers=0 to avoid /dev/shm exhaustion: with the 21 GB train cache,
# PyTorch's DataLoader IPC blows past container /dev/shm limits even at
# n_jobs=1, and `set_sharing_strategy("file_system")` doesn't help because
# _share_filename_cpu_ still routes through shm_open on Linux. Bottleneck
# is the cAdLIF temporal loop submitting CUDA kernels (GIL-bound), so
# parallel data loading buys little here anyway. Override with --workers N
# if running in a container with a large --shm-size.
num_workers = 0
if _CLI_WORKERS is not None:
num_workers = _CLI_WORKERS
net = cAdLIFSSCSNN(
n_channels=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
# under high host-side memory pressure.
train_loader, val_loader, test_loader = setup_dataloaders(
batch_size=batch_size,
num_workers=num_workers,
n_bins=n_bins,
dyned_levels=dyned_levels,
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_snn_v3", 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 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_snn_v3", 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)",
)
args = parser.parse_args()
if args.cpu:
os.environ["CUDA_VISIBLE_DEVICES"] = ""
_CLI_WORKERS = args.workers
if args.gen_data:
import multiprocessing
def _build_with_prefix(prefix, cls, kwargs):
"""Build cache with prefixed output lines."""
import sys
_orig_write = sys.stdout.write
_orig_flush = sys.stdout.flush
def _prefixed_write(s):
if s.strip():
_orig_write(f" [{prefix}] {s}")
else:
_orig_write(s)
_orig_flush()
sys.stdout.write = _prefixed_write
try:
cls(**kwargs)
finally:
sys.stdout.write = _orig_write
combos = []
for levels in [16, 32, 64, 128, 256]:
for n_bins in [25, 50, 100]:
for split in ["train", "valid", "test"]:
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 smaller hosts; streaming keeps the build to the
# spike_trains_arr footprint 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)