Skip to content

Add C weight initializers, LR schedules, and extra losses - #52

Merged
w4ffl35 merged 1 commit into
masterfrom
fix/nir-init-loss-sched
Oct 14, 2026
Merged

w4ffl35 merged 1 commit into
masterfrom
fix/nir-init-loss-sched

Conversation

@w4ffl35

@w4ffl35 w4ffl35 commented Oct 14, 2026

Copy link
Copy Markdown
Member

Closes #45, closes #46, closes #47.

What

New training-side modules (headers in include/nir_c_runtime/, sources in src/, registered in both CMake files). No changes to backward ops, STDP, driver, or checkpoint-state code.

#45 — weight initializers (init.h/init.c)

  • nir_init_uniform, nir_init_xavier_uniform/_normal, nir_init_he_uniform/_normal (torch kaiming fan_in mode), all drawing from nir_rng_t.
  • nir_init_fans implements torch _calculate_fan_in_and_fan_out for Linear/Affine [out, in] and Conv [out, in, *kernel] shapes.
  • Formulas and fan conventions documented in the header.

#46 — LR schedules (schedule.h/schedule.c)

  • Stateless nir_lr_step_decay, nir_lr_cosine, nir_lr_exponential mirroring torch StepLR / CosineAnnealingLR / ExponentialLR with last_epoch = step.
  • Step 0 yields base_lr; the counter is the checkpoint step (completed optimizer steps), so resume continues the curve. Header shows the per-epoch opt->lr write for both SGD and Adam.
  • Cosine follows torch past t_max (rises again, no clamp) — documented.

#47 — losses (loss.h/loss.c)

  • nir_mse_loss: mean-reduction MSE + gradient, mirroring F.mse_loss on one sample; serves readout/membrane regression recipes.
  • nir_spike_count_loss: the decided form is cross-entropy over spike counts (not MSE on counts), matching the snnTorch sum-over-time-then-CE recipe and spikeforge's averaged-readout-plus-cross_entropy path. Caller sums the T-step window; each step's spikes share their class's count gradient.
  • Both keep the single-sample-first 1-D convention; batch reduction stays with Mini-batch training (gradient accumulation, batched loss reduction, seeded shuffling) #42.

Tests

  • test_init: fan values incl. Conv shapes, mean/variance within tolerance at N=20000, seeded bit-identical reproducibility, error shapes.
  • test_init_graph: Xavier-init fc1 weights land via nir_graph_set_weights (read back byte-identical) and the graph runs forward cleanly.
  • test_schedule: closed-form values incl. boundary steps (9/10, 19/20, T/2, T, 2T), SGD + Adam composition through lr, error shapes.
  • test_loss: closed-form MSE and spike-count value + gradient, spike-count bit-equals nir_cross_entropy on the same vector, aliasing rules.
  • Local: 39/39 ctest green in Debug, Release, and ASan/UBSan builds, no compiler warnings under -Wall -Wextra -Werror.

…closes #46, closes #47)

init.h/init.c: torch nn.init initializers (uniform, xavier uniform/normal,
he uniform/normal with fan_in mode) over nir_rng_t, plus a torch fan-in/out
helper for Linear/Affine/Conv weight shapes.

schedule.h/schedule.c: stateless step-decay, cosine-annealing, and
exponential LR schedules mirroring torch StepLR/CosineAnnealingLR/
ExponentialLR, keyed by the checkpoint step counter and applied through
the SGD/Adam lr fields.

loss.h/loss.c: mean-reduction MSE and softmax cross-entropy over spike
counts (the decided spike-count form, matching the snnTorch sum-then-CE
recipe and spikeforge's averaged-readout-plus-CE path).

Tests: test_init (fans, stats, reproducibility, errors), test_init_graph
(set_weights + forward), test_schedule (closed forms, boundaries,
optimizer composition), test_loss (closed forms, CE equivalence).
39/39 ctest green in Debug, Release, and ASan/UBSan.
@w4ffl35
w4ffl35 force-pushed the fix/nir-init-loss-sched branch from 0f3f5d3 to 85007d6 Compare October 14, 2026 11:21
@w4ffl35
w4ffl35 merged commit ba02f38 into master Oct 14, 2026
3 checks passed
@w4ffl35
w4ffl35 deleted the fix/nir-init-loss-sched branch October 14, 2026 11:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

1 participant