Your First Model

import torch
import thrindex as thx
import thrindex.snn as snn
from thrindex.train import rate_loss

model = snn.Sequential(
    snn.Dense(20, 64),
    snn.LIF(tau_mem=20.0, threshold=0.3, reset="subtract"),
    snn.Dense(64, 5),
    snn.LIF(tau_mem=20.0, threshold=0.3, reset="subtract"),
)

import torch

T = 50
n_in = 20
n_classes = 5
n_samples = 500

torch.manual_seed(42)
rates = torch.linspace(0.05, 0.40, n_classes)

def make_batch(n, seed):
    torch.manual_seed(seed)
    labels = torch.randint(0, n_classes, (n,))
    spikes = torch.zeros(n, T, n_in)
    for i, c in enumerate(labels):
        rate = rates[c]
        spikes[i] = (torch.rand(T, n_in) < rate).float()
    return spikes, labels

x_train, y_train = make_batch(n_samples, seed=0)
x_val, y_val = make_batch(100, seed=1)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(10):
    model.train()
    for i in range(0, n_samples, 32):
        x = x_train[i:i+32]
        y = y_train[i:i+32]
        out = model(x)                                         # [batch, T, 5]
        loss = rate_loss(out, y)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
        optimizer.step()
        optimizer.zero_grad()

    model.eval()
    with torch.no_grad():
        val_out = model(x_val)
        counts = val_out.sum(dim=1)                            # [batch, 5]
        preds = counts.argmax(dim=-1)
        acc = (preds == y_val).float().mean().item()
    print(f"Epoch {epoch+1:2d}  val_acc={acc:.3f}")

thx.compile(model, "first_model.thx")
print("Compiled to first_model.thx")

thrindex run first_model.thx

echo '{"spikes": [[0.0, ...]]}' | thrindex run first_model.thx --stdin

thrindex doctor --check first_model.thx

artifact:   first_model.thx
format:     m3
crc32:      OK  (d3a4f82c)
target:     sim
layers:     Dense(20→64) LIF  Dense(64→5) LIF