shd_train.py

"""Surrogate gradient SNN training for the SHD benchmark.

Trains a recurrent SNN (700 -> hidden -> 20) using backpropagation through
time with a fast-sigmoid surrogate gradient.

Supports two neuron models:
  - LIF: multiplicative decay (v = beta * v + (1-beta) * I). Default.
  - adLIF: Adaptive LIF with Symplectic Euler discretization.
    Updates adaptation BEFORE threshold computation for richer temporal dynamics.
    Published: 95.81% on SHD (SE-adLIF, 2025).

Hardware mapping (CUBA neuron, P22A):
    decay_u = round(alpha * 4096)   (12-bit fractional)

Usage:
    python shd_train.py --data-dir data/shd --epochs 200 --hidden 512
    python shd_train.py --neuron-type adlif --dropout 0.15 --epochs 200
"""

import os
import sys
import random
import argparse
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

# Add benchmarks dir to path for shd_loader import
sys.path.insert(0, os.path.dirname(__file__))
from shd_loader import SHDDataset, collate_fn, N_CHANNELS, N_CLASSES


# ---------------------------------------------------------------------------
# Surrogate gradient
# ---------------------------------------------------------------------------

class SurrogateSpikeFunction(torch.autograd.Function):
    """Heaviside forward, fast-sigmoid backward (surrogate gradient)."""

    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return (x >= 0).float()

    @staticmethod
    def backward(ctx, grad_output):
        x, = ctx.saved_tensors
        # Fast sigmoid surrogate: 1 / (1 + scale*|x|)^2
        scale = 25.0
        grad = grad_output / (scale * torch.abs(x) + 1.0) ** 2
        return grad


surrogate_spike = SurrogateSpikeFunction.apply


# ---------------------------------------------------------------------------
# Neuron model — multiplicative decay LIF (maps to CUBA hardware neuron)
# ---------------------------------------------------------------------------

class LIFNeuron(nn.Module):
    """Leaky Integrate-and-Fire with multiplicative (exponential) decay.

    Dynamics per timestep:
        v = beta * v_prev + (1 - beta) * I   # exponential decay + scaled input
        spike = Heaviside(v - threshold)       # surrogate in backward
        v = v * (1 - spike)                    # hard reset

    Hardware mapping (CUBA neuron, P22A):
        decay_u = round(beta * 4096)   (12-bit fractional)
    """

    def __init__(self, size, beta_init=0.95, threshold=1.0, learn_beta=True):
        super().__init__()
        self.size = size
        self.threshold = threshold
        # Learnable time constant via sigmoid-mapped beta
        if learn_beta:
            # Initialize so sigmoid(x) = beta_init
            init_val = np.log(beta_init / (1.0 - beta_init))
            self.beta_raw = nn.Parameter(torch.full((size,), init_val))
        else:
            self.register_buffer('beta_raw',
                                 torch.full((size,), np.log(beta_init / (1.0 - beta_init))))

    @property
    def beta(self):
        return torch.sigmoid(self.beta_raw)

    def forward(self, input_current, v_prev):
        beta = self.beta
        v = beta * v_prev + (1.0 - beta) * input_current
        spikes = surrogate_spike(v - self.threshold)
        v = v * (1.0 - spikes)   # hard reset to 0
        return v, spikes


# ---------------------------------------------------------------------------
# Adaptive LIF neuron — Symplectic Euler discretization
# ---------------------------------------------------------------------------

class AdaptiveLIFNeuron(nn.Module):
    """Adaptive LIF with Symplectic Euler (SE) discretization.

    Key: adaptation is updated BEFORE threshold computation, so the neuron
    can anticipate its own spike — greatly improves temporal coding.

    Dynamics per timestep (SE order):
        a = rho * a_prev + spike_prev          # 1. adaptation update FIRST
        theta = threshold_base + beta_a * a    # 2. adaptive threshold
        v = alpha * v_prev + (1-alpha) * I     # 3. membrane update
        spike = Heaviside(v - theta)            # 4. spike decision
        v = v * (1 - spike)                     # 5. hard reset

    Hardware note: adaptation is training-only. Only alpha (membrane decay)
    deploys to CUBA hardware as decay_v = round(alpha * 4096).
    """

    def __init__(self, size, alpha_init=0.90, rho_init=0.85, beta_a_init=1.8,
                 threshold=1.0):
        super().__init__()
        self.size = size
        self.threshold_base = nn.Parameter(torch.full((size,), threshold))

        # Membrane decay (learnable via sigmoid)
        init_alpha = np.log(alpha_init / (1.0 - alpha_init))
        self.alpha_raw = nn.Parameter(torch.full((size,), init_alpha))

        # Adaptation decay (learnable via sigmoid)
        init_rho = np.log(rho_init / (1.0 - rho_init))
        self.rho_raw = nn.Parameter(torch.full((size,), init_rho))

        # Adaptation strength (learnable, softplus to keep positive)
        # softplus^{-1}(beta_a_init) = log(exp(beta_a_init) - 1)
        init_beta_a = np.log(np.exp(beta_a_init) - 1.0)
        self.beta_a_raw = nn.Parameter(torch.full((size,), init_beta_a))

    @property
    def alpha(self):
        return torch.sigmoid(self.alpha_raw)

    def forward(self, input_current, v_prev, a_prev, spike_prev):
        alpha = torch.sigmoid(self.alpha_raw)
        rho = torch.sigmoid(self.rho_raw)
        beta_a = F.softplus(self.beta_a_raw)

        # SE discretization: adaptation FIRST
        a_new = rho * a_prev + spike_prev
        theta = self.threshold_base + beta_a * a_new

        # Membrane dynamics
        v = alpha * v_prev + (1.0 - alpha) * input_current
        spikes = surrogate_spike(v - theta)
        v = v * (1.0 - spikes)  # hard reset

        return v, spikes, a_new


# ---------------------------------------------------------------------------
# Event-drop data augmentation
# ---------------------------------------------------------------------------

def event_drop_augment(spikes_batch, drop_time_prob=0.1, drop_neuron_prob=0.05):
    """Randomly drop entire time bins or channels for regularization.

    Operates on full batch (B, T, C) for efficiency. ~1% accuracy boost.
    """
    if random.random() < 0.5:
        # Drop-by-time: zero out random time bins (shared across batch)
        B, T, C = spikes_batch.shape
        mask = (torch.rand(1, T, 1, device=spikes_batch.device)
                > drop_time_prob).float()
        return spikes_batch * mask
    else:
        # Drop-by-neuron: zero out random input channels (shared across batch)
        B, T, C = spikes_batch.shape
        mask = (torch.rand(1, 1, C, device=spikes_batch.device)
                > drop_neuron_prob).float()
        return spikes_batch * mask


# ---------------------------------------------------------------------------
# SNN model
# ---------------------------------------------------------------------------

class SHDSNN(nn.Module):
    """Recurrent SNN for SHD classification.

    700 (input spikes) -> hidden (recurrent LIF/adLIF) -> 20 (non-spiking readout)
    Readout: time-summed membrane potential of output layer -> softmax.
    """

    def __init__(self, n_input=N_CHANNELS, n_hidden=256, n_output=N_CLASSES,
                 beta_hidden=0.95, beta_out=0.9, threshold=1.0, dropout=0.3,
                 neuron_type='lif', alpha_init=0.90, rho_init=0.85,
                 beta_a_init=1.8):
        super().__init__()
        self.n_hidden = n_hidden
        self.n_output = n_output
        self.dropout_p = dropout
        self.neuron_type = neuron_type

        # Synaptic weight matrices
        self.fc1 = nn.Linear(n_input, n_hidden, bias=False)
        self.fc2 = nn.Linear(n_hidden, n_output, bias=False)

        # Recurrent connection in hidden layer
        self.fc_rec = nn.Linear(n_hidden, n_hidden, bias=False)

        # Hidden layer neuron
        if neuron_type == 'adlif':
            self.lif1 = AdaptiveLIFNeuron(
                n_hidden, alpha_init=alpha_init, rho_init=rho_init,
                beta_a_init=beta_a_init, threshold=threshold)
        else:
            self.lif1 = LIFNeuron(n_hidden, beta_init=beta_hidden,
                                   threshold=threshold, learn_beta=True)

        # Output layer always standard LIF (readout doesn't need adaptation)
        self.lif2 = LIFNeuron(n_output, beta_init=beta_out,
                               threshold=threshold, learn_beta=True)

        # Dropout for regularization
        self.dropout = nn.Dropout(p=dropout)

        # Weight init
        nn.init.xavier_uniform_(self.fc1.weight, gain=0.5)
        nn.init.xavier_uniform_(self.fc2.weight, gain=0.5)
        nn.init.orthogonal_(self.fc_rec.weight, gain=0.2)

    def forward(self, x):
        """Forward pass unrolled through T timesteps.

        Args:
            x: (batch, T, n_input) dense spike input

        Returns:
            output: (batch, n_output) averaged membrane for classification
        """
        batch, T, _ = x.shape
        device = x.device

        v1 = torch.zeros(batch, self.n_hidden, device=device)
        v2 = torch.zeros(batch, self.n_output, device=device)
        spk1 = torch.zeros(batch, self.n_hidden, device=device)

        out_sum = torch.zeros(batch, self.n_output, device=device)

        # adLIF needs adaptation state
        if self.neuron_type == 'adlif':
            a1 = torch.zeros(batch, self.n_hidden, device=device)

        for t in range(T):
            # Hidden layer: feedforward + recurrent
            I1 = self.fc1(x[:, t]) + self.fc_rec(spk1)

            if self.neuron_type == 'adlif':
                v1, spk1, a1 = self.lif1(I1, v1, a1, spk1)
            else:
                v1, spk1 = self.lif1(I1, v1)

            # Apply dropout to hidden spikes
            spk1_drop = self.dropout(spk1) if self.training else spk1

            # Output layer (non-spiking readout: integrate with decay)
            I2 = self.fc2(spk1_drop)
            beta_out = self.lif2.beta
            v2 = beta_out * v2 + (1.0 - beta_out) * I2

            out_sum = out_sum + v2

        # Normalize by timesteps
        return out_sum / T


# ---------------------------------------------------------------------------
# Training loop
# ---------------------------------------------------------------------------

def train_epoch(model, loader, optimizer, device, use_event_drop=False,
                label_smoothing=0.0):
    model.train()
    total_loss = 0.0
    correct = 0
    total = 0

    for inputs, labels in loader:
        inputs, labels = inputs.to(device), labels.to(device)

        # Event-drop augmentation (batch-level for efficiency)
        if use_event_drop:
            inputs = event_drop_augment(inputs)

        optimizer.zero_grad()
        output = model(inputs)
        loss = F.cross_entropy(output, labels, label_smoothing=label_smoothing)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()

        total_loss += loss.item() * inputs.size(0)
        correct += (output.argmax(1) == labels).sum().item()
        total += inputs.size(0)

    return total_loss / total, correct / total


@torch.no_grad()
def evaluate(model, loader, device):
    model.eval()
    total_loss = 0.0
    correct = 0
    total = 0

    for inputs, labels in loader:
        inputs, labels = inputs.to(device), labels.to(device)

        output = model(inputs)
        loss = F.cross_entropy(output, labels)

        total_loss += loss.item() * inputs.size(0)
        correct += (output.argmax(1) == labels).sum().item()
        total += inputs.size(0)

    return total_loss / total, correct / total


def main():
    parser = argparse.ArgumentParser(description="Train SNN on SHD benchmark")
    parser.add_argument("--data-dir", default="data/shd")
    parser.add_argument("--epochs", type=int, default=200)
    parser.add_argument("--batch-size", type=int, default=128)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--weight-decay", type=float, default=1e-4)
    parser.add_argument("--hidden", type=int, default=512)
    parser.add_argument("--threshold", type=float, default=1.0)
    parser.add_argument("--beta-hidden", type=float, default=0.95,
                        help="Initial membrane decay factor for hidden layer")
    parser.add_argument("--beta-out", type=float, default=0.9,
                        help="Initial membrane decay factor for output layer")
    parser.add_argument("--dropout", type=float, default=0.3)
    parser.add_argument("--dt", type=float, default=4e-3,
                        help="Time bin width in seconds (4ms -> 250 bins)")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--save", default="shd_model.pt")
    parser.add_argument("--no-recurrent", action="store_true",
                        help="Disable recurrent hidden connection")
    parser.add_argument("--neuron-type", choices=["lif", "adlif"], default="lif",
                        help="Neuron model: lif (standard) or adlif (adaptive, SE)")
    parser.add_argument("--alpha-init", type=float, default=0.90,
                        help="Initial membrane decay for adLIF (default: 0.90)")
    parser.add_argument("--rho-init", type=float, default=0.85,
                        help="Initial adaptation decay for adLIF (default: 0.85)")
    parser.add_argument("--beta-a-init", type=float, default=1.8,
                        help="Initial adaptation strength for adLIF (default: 1.8)")
    parser.add_argument("--event-drop", action="store_true", default=None,
                        help="Enable event-drop augmentation (auto-enabled for adlif)")
    parser.add_argument("--label-smoothing", type=float, default=0.0,
                        help="Label smoothing factor (0.0=off, 0.1=recommended)")
    args = parser.parse_args()

    # Auto-enable event-drop for adLIF if not explicitly set
    if args.event_drop is None:
        args.event_drop = (args.neuron_type == 'adlif')

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    # Dataset
    print("Loading SHD dataset...")
    train_ds = SHDDataset(args.data_dir, "train", dt=args.dt)
    test_ds = SHDDataset(args.data_dir, "test", dt=args.dt)

    train_loader = DataLoader(
        train_ds, batch_size=args.batch_size, shuffle=True,
        collate_fn=collate_fn, num_workers=0, pin_memory=True)
    test_loader = DataLoader(
        test_ds, batch_size=args.batch_size, shuffle=False,
        collate_fn=collate_fn, num_workers=0, pin_memory=True)

    print(f"Train: {len(train_ds)}, Test: {len(test_ds)}, "
          f"Time bins: {train_ds.n_bins} (dt={args.dt*1000:.1f}ms)")

    # Model
    model = SHDSNN(
        n_hidden=args.hidden,
        threshold=args.threshold,
        beta_hidden=args.beta_hidden,
        beta_out=args.beta_out,
        dropout=args.dropout,
        neuron_type=args.neuron_type,
        alpha_init=args.alpha_init,
        rho_init=args.rho_init,
        beta_a_init=args.beta_a_init,
    ).to(device)

    if args.no_recurrent:
        model.fc_rec.weight.data.zero_()
        model.fc_rec.weight.requires_grad = False

    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    neuron_info = args.neuron_type.upper()
    if args.neuron_type == 'adlif':
        neuron_info += f" (alpha={args.alpha_init}, rho={args.rho_init}, beta_a={args.beta_a_init})"
    print(f"Model: {N_CHANNELS}->{args.hidden}->{N_CLASSES}, "
          f"{n_params:,} params ({neuron_info}, "
          f"recurrent={'off' if args.no_recurrent else 'on'}, "
          f"dropout={args.dropout}, event_drop={args.event_drop})")

    # Optimizer with weight decay
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr,
                                   weight_decay=args.weight_decay)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.epochs,
                                                            eta_min=1e-5)

    best_acc = 0.0
    for epoch in range(args.epochs):
        train_loss, train_acc = train_epoch(model, train_loader, optimizer, device,
                                               use_event_drop=args.event_drop,
                                               label_smoothing=args.label_smoothing)
        test_loss, test_acc = evaluate(model, test_loader, device)
        scheduler.step()

        if test_acc > best_acc:
            best_acc = test_acc
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'test_acc': test_acc,
                'args': vars(args),
            }, args.save)

        lr = optimizer.param_groups[0]['lr']
        print(f"Epoch {epoch+1:3d}/{args.epochs} | "
              f"Train: {train_loss:.4f} / {train_acc*100:.1f}% | "
              f"Test: {test_loss:.4f} / {test_acc*100:.1f}% | "
              f"LR={lr:.2e} | Best={best_acc*100:.1f}%")

    print(f"\nDone. Best test accuracy: {best_acc*100:.1f}%")
    print(f"Model saved to {args.save}")


if __name__ == "__main__":
    main()