From a6bee2b29a96afac56ea251ace4420f9ff29df6c Mon Sep 17 00:00:00 2001 From: Joe Date: Tue, 13 Oct 2026 19:42:25 -0600 Subject: [PATCH] Add C-side training math: surrogate hook, backward ops, loss, optimizers Forward spikes stay strict Heavisides; the backward pass substitutes a named surrogate derivative (fast_sigmoid/atan, snnTorch formulas) with no silent default. New train.h/train.c provide name resolution, softmax cross-entropy, SGD/Adam over raw weight buffers, and a local STDP rule ported from STDPSynapse with decay-before-update ordering. Layer backward ops cover Threshold, LI, Affine, and Linear; test_train checks each primitive plus one Adam step on a two-layer topology against float64 closed-form references and a 300-step convergence run. Also fix two Conv2d groups set_error call sites left with a stale third argument after #24 (same fix as the loader PR; master did not compile). A graph-level backprop-through-time tape and a committed torch fixture generator remain follow-ups. Fixes #26. --- CMakeLists.txt | 1 + README.md | 10 +- include/nir_c_runtime/ops.h | 61 ++++ include/nir_c_runtime/train.h | 133 +++++++++ src/graph.c | 6 +- src/ops.c | 192 +++++++++++++ src/train.c | 332 +++++++++++++++++++++ tests/CMakeLists.txt | 1 + tests/test_train.c | 526 ++++++++++++++++++++++++++++++++++ 9 files changed, 1259 insertions(+), 3 deletions(-) create mode 100644 include/nir_c_runtime/train.h create mode 100644 src/train.c create mode 100644 tests/test_train.c diff --git a/CMakeLists.txt b/CMakeLists.txt index a335e79..5f8b125 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -17,6 +17,7 @@ add_library(nir_c_runtime src/ops.c src/runtime.c src/fixed_event.c + src/train.c ) target_include_directories(nir_c_runtime PUBLIC include) target_link_libraries(nir_c_runtime PUBLIC m) diff --git a/README.md b/README.md index 2cff9e4..a8fdc87 100644 --- a/README.md +++ b/README.md @@ -74,10 +74,18 @@ target-specific kernels remain future optimization work. - [x] Allocation-free fixed-point event-module ABI for bounded, single-output temporal detectors (`fixed_event.h`), deterministic guard-state/action helpers, and an 8,192-row exhaustive C/Curlee truth-table fingerprint. +- [x] C-side training math (`train.h`): named surrogate gradients + (fast_sigmoid/atan, no silent default) with Threshold/LI/Affine/Linear + backward ops, softmax cross-entropy, SGD/Adam optimizers over raw + weight buffers, and a local STDP rule. Dataset loading, RNG/encoding, + checkpoints, and a graph-level backprop-through-time tape remain + future work (see the training epic). ## Out of scope (for now) -- Training. This is inference-only. +- Full training pipelines. The loss/optimizer/STDP primitives above are + supported; data loading, spike encoding, checkpointing, and graph-level + backprop-through-time are not built yet. - General Curlee/JOE graph execution. The narrow, allocation-free event ABI is integrated as the reference path for [curlee#316](https://github.com/w4ffl35/curlee/issues/316); loading arbitrary diff --git a/include/nir_c_runtime/ops.h b/include/nir_c_runtime/ops.h index 36ca47b..010a027 100644 --- a/include/nir_c_runtime/ops.h +++ b/include/nir_c_runtime/ops.h @@ -166,4 +166,65 @@ int nir_op_scale_into( nir_tensor_t *output ); +/* Training: single-step backward ops. The forward spike + * nonlinearities stay strict Heavisides; the backward pass below + * substitutes a smooth surrogate derivative, mirroring snnTorch's + * `surrogate` module (see spikeforge/introspection/surrogate.py). + * Names resolve through nir_surrogate_from_name in train.h, which + * rejects unknown names instead of silently defaulting. */ + +typedef enum { + NIR_SURROGATE_FAST_SIGMOID, + NIR_SURROGATE_ATAN +} nir_surrogate_kind_t; + +/* Surrogate derivative dS/dx at x. fast_sigmoid with slope s: + * s / (1 + |s*x|)^2 (snnTorch default s = 25). atan with + * steepness a: (a/2) / (1 + (pi*a*x/2)^2) (snnTorch default + * a = 2). Returns 0 unless `kind` is known and `slope` is + * finite and positive. */ +int nir_surrogate_grad( + nir_surrogate_kind_t kind, float slope, float x, float *out +); + +/* Backward through Threshold: forward is the strict Heaviside, + * and grad_input[i] = grad_output[i] * surr(membrane[i] - + * threshold). `membrane` is the retained pre-threshold input. + * Shapes must match; `grad_input` may exactly alias + * `grad_output` (elementwise rule). */ +int nir_op_threshold_backward( + const nir_threshold_params_t *params, const nir_tensor_t *membrane, + const nir_tensor_t *grad_output, nir_surrogate_kind_t surrogate, + float slope, nir_tensor_t *grad_input +); + +/* Backward through one LI step: with decay = exp(-1 / tau), + * grad_input = grad_output * r * (1 - decay) and grad_prev = + * grad_output * decay. `grad_prev` may be NULL when the + * previous-state gradient is unneeded. Either output may exactly + * alias `grad_output`; the two outputs must not overlap each + * other. */ +int nir_op_li_backward( + const nir_li_params_t *params, const nir_tensor_t *grad_output, + nir_tensor_t *grad_input, nir_tensor_t *grad_prev +); + +/* Backward through y = Wx + b: grad_weight += outer(grad_output, + * input), grad_bias += grad_output, grad_input = W^T grad_output. + * Weight and bias gradients accumulate (the caller zeroes them); + * `grad_bias` may be NULL when the node has no bias. All three + * outputs must be disjoint from the inputs and from each other. */ +int nir_op_affine_backward( + const nir_affine_params_t *params, const nir_tensor_t *input, + const nir_tensor_t *grad_output, nir_tensor_t *grad_weight, + nir_tensor_t *grad_bias, nir_tensor_t *grad_input +); + +/* Backward through y = Wx: as above without the bias term. */ +int nir_op_linear_backward( + const nir_linear_params_t *params, const nir_tensor_t *input, + const nir_tensor_t *grad_output, nir_tensor_t *grad_weight, + nir_tensor_t *grad_input +); + #endif /* NIR_C_RUNTIME_OPS_H */ diff --git a/include/nir_c_runtime/train.h b/include/nir_c_runtime/train.h new file mode 100644 index 0000000..95088fe --- /dev/null +++ b/include/nir_c_runtime/train.h @@ -0,0 +1,133 @@ +/* C-side training math: surrogate-name resolution, cross-entropy + * loss, SGD/Adam optimizers over raw weight buffers, and a local + * STDP rule ported from spikeforge's STDPSynapse. + * + * This mirrors the torch training step (cross_entropy on logits, + * then a torch.optim.Adam step) for single-step topologies whose + * backward pass the caller composes from the layer backward ops in + * ops.h. A graph-level tape retaining per-step activations for + * backprop-through-time is future work; optimizers here operate on + * caller-provided float buffers (e.g. a node's weight array). + */ +#ifndef NIR_C_RUNTIME_TRAIN_H +#define NIR_C_RUNTIME_TRAIN_H + +#include + +#include "nir_c_runtime/ops.h" +#include "nir_c_runtime/tensor.h" + +/* Resolve "fast_sigmoid" or "atan" to a surrogate kind. Any other + * name (or NULL) fails with a malloc'd error naming it, mirroring + * spikeforge's resolve_surrogate raising KeyError -- there is no + * silent default surrogate. */ +int nir_surrogate_from_name( + const char *name, nir_surrogate_kind_t *out, char **error +); + +/* Softmax cross-entropy for one logit vector, matching + * torch.nn.functional.cross_entropy on a single sample: loss = + * -log(softmax(logits)[target]), grad = softmax(logits) with + * grad[target] -= 1. Batch reduction is the caller's job. */ +int nir_cross_entropy( + const nir_tensor_t *logits, size_t target, float *loss, + nir_tensor_t *grad +); + +/* SGD with optional momentum and L2 weight decay, matching + * torch.optim.SGD with dampening=0 and nesterov=False: + * d = grad + weight_decay * w + * v = momentum * v + d + * w -= lr * v + */ +typedef struct { + float lr; + float momentum; + float weight_decay; + float *velocity; + size_t count; +} nir_sgd_t; + +/* lr must be finite and >= 0, momentum and weight_decay finite + * and >= 0, count nonzero. The velocity buffer starts at zero. */ +int nir_sgd_init( + nir_sgd_t *opt, size_t count, float lr, float momentum, + float weight_decay, char **error +); +int nir_sgd_step( + nir_sgd_t *opt, float *weights, const float *grad, size_t count +); +void nir_sgd_free(nir_sgd_t *opt); + +/* Adam with bias correction, matching torch.optim.Adam with + * amsgrad=False: + * d = grad + weight_decay * w + * m = beta1 * m + (1 - beta1) * d + * v = beta2 * v + (1 - beta2) * d^2 + * w -= lr * (m / (1 - beta1^t)) / + * (sqrt(v / (1 - beta2^t)) + eps) + * where t counts steps from 1. + */ +typedef struct { + float lr; + float beta1; + float beta2; + float eps; + float weight_decay; + float *m; + float *v; + size_t count; + unsigned long step; +} nir_adam_t; + +/* lr must be finite and >= 0, betas finite in [0, 1), eps finite + * and > 0, weight_decay finite and >= 0, count nonzero. */ +int nir_adam_init( + nir_adam_t *opt, size_t count, float lr, float beta1, + float beta2, float eps, float weight_decay, char **error +); +int nir_adam_step( + nir_adam_t *opt, float *weights, const float *grad, size_t count +); +void nir_adam_free(nir_adam_t *opt); + +/* Pair-based STDP with pre/post eligibility traces, ported from + * spikeforge.memory.STDPSynapse (Song/Abbott rule). Weights are + * [in_features][out_features], row-major pre-major, matching the + * torch [in, out] layout and the spikes @ weight forward. */ +typedef struct { + size_t in_features; + size_t out_features; + float tau_pre; + float tau_post; + float a_plus; + float a_minus; + float w_min; + float w_max; + float *weight; + float *pre_trace; + float *post_trace; +} nir_stdp_t; + +/* taus must be finite and > 0, w_min <= w_max, both sizes nonzero. + * Weights and traces start at zero. */ +int nir_stdp_init( + nir_stdp_t *s, size_t in_features, size_t out_features, + float tau_pre, float tau_post, float a_plus, float a_minus, + float w_min, float w_max, char **error +); + +/* Overwrite the weight matrix (in*out floats, pre-major). */ +int nir_stdp_set_weight(nir_stdp_t *s, const float *weight); + +/* One STDP step: decay both traces first (a spike reads the + * *prior* step's trace), potentiate on post spikes landing on + * lingering pre trace, depress on pre spikes landing on lingering + * post trace, clamp to [w_min, w_max], then add this step's + * spikes to the traces. */ +int nir_stdp_step( + nir_stdp_t *s, const float *pre_spikes, const float *post_spikes +); +void nir_stdp_free(nir_stdp_t *s); + +#endif /* NIR_C_RUNTIME_TRAIN_H */ diff --git a/src/graph.c b/src/graph.c index 8cee597..61d685b 100644 --- a/src/graph.c +++ b/src/graph.c @@ -324,11 +324,13 @@ static int load_conv2d( int groups = 1; const json_value_t *groups_json = json_object_get(params, "groups"); if (groups_json != NULL && !int_at(params, "groups", &groups)) { - set_error(error, "Conv2d groups must be a nonnegative integer", NULL); + set_error(error, "Conv2d groups must be a nonnegative integer"); return 0; } if (groups != 1) { - set_error(error, "Conv2d groups must be 1; grouped convolution is not supported", NULL); + set_error( + error, "Conv2d groups must be 1; grouped convolution is not supported" + ); return 0; } if (!pair_at(params, "stride", &p->stride_h, &p->stride_w) || diff --git a/src/ops.c b/src/ops.c index 34c256f..6daef6b 100644 --- a/src/ops.c +++ b/src/ops.c @@ -722,3 +722,195 @@ nir_tensor_t nir_op_scale( } return output; } + +static double surrogate_grad_unchecked( + nir_surrogate_kind_t kind, double slope, double x +) { + if (kind == NIR_SURROGATE_FAST_SIGMOID) { + double d = 1.0 + fabs(slope * x); + return slope / (d * d); + } + double d = 1.5707963267948966 * slope * x; + return (slope / 2.0) / (1.0 + d * d); +} + +int nir_surrogate_grad( + nir_surrogate_kind_t kind, float slope, float x, float *out +) { + if (out == NULL || !isfinite(slope) || slope <= 0.0f || isnan(x)) { + return 0; + } + if (kind != NIR_SURROGATE_FAST_SIGMOID && + kind != NIR_SURROGATE_ATAN) { + return 0; + } + *out = (float)surrogate_grad_unchecked( + kind, (double)slope, (double)x + ); + return 1; +} + +int nir_op_threshold_backward( + const nir_threshold_params_t *params, const nir_tensor_t *membrane, + const nir_tensor_t *grad_output, nir_surrogate_kind_t surrogate, + float slope, nir_tensor_t *grad_input +) { + float probe; + if (params == NULL || !isfinite(params->threshold) || + grad_input == NULL || + !tensor_valid( + membrane, membrane != NULL ? membrane->ndim : 0) || + !tensor_valid( + grad_output, membrane != NULL ? membrane->ndim : 0) || + !tensor_same_shape(membrane, grad_output) || + !nir_surrogate_grad(surrogate, slope, 0.0f, &probe)) { + return 0; + } + if (tensors_partially_overlap(membrane, grad_input) || + tensors_partially_overlap(grad_output, grad_input) || + !prepare_output(grad_input, membrane->shape, membrane->ndim)) { + return 0; + } + double threshold = (double)params->threshold; + double s = (double)slope; + for (size_t i = 0; i < membrane->size; i++) { + double g = surrogate_grad_unchecked( + surrogate, s, (double)membrane->data[i] - threshold + ); + grad_input->data[i] = grad_output->data[i] * (float)g; + } + return 1; +} + +int nir_op_li_backward( + const nir_li_params_t *params, const nir_tensor_t *grad_output, + nir_tensor_t *grad_input, nir_tensor_t *grad_prev +) { + if (params == NULL || grad_input == NULL || + !tensor_valid( + grad_output, grad_output != NULL ? grad_output->ndim : 0) || + !isfinite(params->tau) || params->tau <= 0.0f || + !isfinite(params->r)) { + return 0; + } + if (tensors_partially_overlap(grad_output, grad_input) || + !prepare_output( + grad_input, grad_output->shape, grad_output->ndim)) { + return 0; + } + if (grad_prev != NULL) { + if (tensors_overlap(grad_input, grad_prev) || + tensors_partially_overlap(grad_output, grad_prev) || + !prepare_output( + grad_prev, grad_output->shape, grad_output->ndim)) { + return 0; + } + } + /* Matches nir_op_li_into's zero-order-hold coefficients. */ + double decay = exp(-1.0 / (double)params->tau); + float to_input = (float)((double)params->r * (1.0 - decay)); + float to_prev = (float)decay; + for (size_t i = 0; i < grad_output->size; i++) { + grad_input->data[i] = grad_output->data[i] * to_input; + } + if (grad_prev != NULL) { + for (size_t i = 0; i < grad_output->size; i++) { + grad_prev->data[i] = grad_output->data[i] * to_prev; + } + } + return 1; +} + +static int affine_backward_into( + const float *weight, const float *bias, size_t out_features, + size_t in_features, const nir_tensor_t *input, + const nir_tensor_t *grad_output, nir_tensor_t *grad_weight, + nir_tensor_t *grad_bias, nir_tensor_t *grad_input +) { + if (weight == NULL || !tensor_valid(input, 1) || + input->size != in_features || + !tensor_valid(grad_output, 1) || + grad_output->size != out_features || grad_weight == NULL || + grad_input == NULL) { + return 0; + } + size_t weight_count; + if (!checked_mul(out_features, in_features, &weight_count)) { + return 0; + } + if ((bias == NULL) != (grad_bias == NULL)) { + return 0; + } + size_t wshape[2] = {out_features, in_features}; + size_t bshape[1] = {out_features}; + size_t xshape[1] = {in_features}; + if (overlaps_output_shape(input, grad_weight, wshape, 2) || + overlaps_output_shape(grad_output, grad_weight, wshape, 2) || + overlaps_output_shape(input, grad_input, xshape, 1) || + overlaps_output_shape(grad_output, grad_input, xshape, 1) || + !prepare_output(grad_weight, wshape, 2) || + !prepare_output(grad_input, xshape, 1)) { + return 0; + } + if (grad_bias != NULL && + (overlaps_output_shape(input, grad_bias, bshape, 1) || + overlaps_output_shape(grad_output, grad_bias, bshape, 1) || + !prepare_output(grad_bias, bshape, 1))) { + return 0; + } + if (tensors_overlap(grad_weight, grad_input) || + (grad_bias != NULL && + (tensors_overlap(grad_weight, grad_bias) || + tensors_overlap(grad_input, grad_bias)))) { + return 0; + } + for (size_t o = 0; o < out_features; o++) { + float gy = grad_output->data[o]; + for (size_t i = 0; i < in_features; i++) { + grad_weight->data[o * in_features + i] += + gy * input->data[i]; + } + if (grad_bias != NULL) { + grad_bias->data[o] += gy; + } + } + for (size_t i = 0; i < in_features; i++) { + float sum = 0.0f; + for (size_t o = 0; o < out_features; o++) { + sum += weight[o * in_features + i] * grad_output->data[o]; + } + grad_input->data[i] = sum; + } + return 1; +} + +int nir_op_affine_backward( + const nir_affine_params_t *params, const nir_tensor_t *input, + const nir_tensor_t *grad_output, nir_tensor_t *grad_weight, + nir_tensor_t *grad_bias, nir_tensor_t *grad_input +) { + if (params == NULL) { + return 0; + } + return affine_backward_into( + params->weight, params->bias, params->out_features, + params->in_features, input, grad_output, grad_weight, + grad_bias, grad_input + ); +} + +int nir_op_linear_backward( + const nir_linear_params_t *params, const nir_tensor_t *input, + const nir_tensor_t *grad_output, nir_tensor_t *grad_weight, + nir_tensor_t *grad_input +) { + if (params == NULL) { + return 0; + } + /* The NIR Linear node carries no bias term. */ + return affine_backward_into( + params->weight, NULL, params->out_features, + params->in_features, input, grad_output, grad_weight, + NULL, grad_input + ); +} diff --git a/src/train.c b/src/train.c new file mode 100644 index 0000000..7ef1fe1 --- /dev/null +++ b/src/train.c @@ -0,0 +1,332 @@ +#include "nir_c_runtime/train.h" + +#include +#include +#include +#include +#include + +static void set_error(char **error, const char *message) { + if (error == NULL) return; + *error = NULL; + if (message == NULL) return; + size_t n = strlen(message) + 1; + char *copy = (char *)malloc(n); + if (copy != NULL) memcpy(copy, message, n); + *error = copy; +} + +static void set_error_arg( + char **error, const char *before, const char *name, + const char *after +) { + if (error == NULL) return; + *error = NULL; + if (before == NULL || name == NULL || after == NULL) return; + size_t n = strlen(before) + strlen(name) + strlen(after) + 1; + char *copy = (char *)malloc(n); + if (copy == NULL) return; + copy[0] = '\0'; + strcat(copy, before); + strcat(copy, name); + strcat(copy, after); + *error = copy; +} + +int nir_surrogate_from_name( + const char *name, nir_surrogate_kind_t *out, char **error +) { + if (error != NULL) *error = NULL; + if (out == NULL) { + set_error(error, "surrogate output must not be null"); + return 0; + } + if (name == NULL) { + set_error(error, "unknown surrogate: (null)"); + return 0; + } + if (strcmp(name, "fast_sigmoid") == 0) { + *out = NIR_SURROGATE_FAST_SIGMOID; + return 1; + } + if (strcmp(name, "atan") == 0) { + *out = NIR_SURROGATE_ATAN; + return 1; + } + set_error_arg(error, "unknown surrogate: ", name, ""); + return 0; +} + +int nir_cross_entropy( + const nir_tensor_t *logits, size_t target, float *loss, + nir_tensor_t *grad +) { + if (logits == NULL || logits->data == NULL || logits->ndim != 1 || + logits->size == 0 || logits->capacity < logits->size || + target >= logits->size || loss == NULL || grad == NULL || + grad->data == NULL || grad->capacity < logits->size) { + return 0; + } + for (size_t i = 0; i < logits->size; i++) { + if (!isfinite(logits->data[i])) return 0; + } + if (logits->data != grad->data) { + uintptr_t lo = (uintptr_t)logits->data; + uintptr_t go = (uintptr_t)grad->data; + size_t bytes = logits->size * sizeof(float); + if (lo < go + bytes && go < lo + bytes) return 0; + } + float peak = logits->data[0]; + for (size_t i = 1; i < logits->size; i++) { + if (logits->data[i] > peak) peak = logits->data[i]; + } + double total = 0.0; + for (size_t i = 0; i < logits->size; i++) { + total += exp((double)logits->data[i] - (double)peak); + } + grad->ndim = 1; + grad->shape[0] = logits->size; + grad->shape[1] = 0; + grad->shape[2] = 0; + grad->size = logits->size; + for (size_t i = 0; i < logits->size; i++) { + double p = exp((double)logits->data[i] - (double)peak) / total; + grad->data[i] = (float)(p - (i == target ? 1.0 : 0.0)); + } + *loss = (float)-log( + exp((double)logits->data[target] - (double)peak) / total + ); + return 1; +} + +int nir_sgd_init( + nir_sgd_t *opt, size_t count, float lr, float momentum, + float weight_decay, char **error +) { + if (error != NULL) *error = NULL; + if (opt == NULL) { + set_error(error, "optimizer output must not be null"); + return 0; + } + memset(opt, 0, sizeof(*opt)); + if (count == 0 || !isfinite(lr) || lr < 0.0f || + !isfinite(momentum) || momentum < 0.0f || + !isfinite(weight_decay) || weight_decay < 0.0f) { + set_error(error, "invalid SGD hyperparameters"); + return 0; + } + if (count > SIZE_MAX / sizeof(float)) { + set_error(error, "optimizer state is too large"); + return 0; + } + opt->velocity = (float *)calloc(count, sizeof(float)); + if (opt->velocity == NULL) { + set_error(error, "out of memory"); + return 0; + } + opt->lr = lr; + opt->momentum = momentum; + opt->weight_decay = weight_decay; + opt->count = count; + return 1; +} + +int nir_sgd_step( + nir_sgd_t *opt, float *weights, const float *grad, size_t count +) { + if (opt == NULL || opt->velocity == NULL || + opt->count != count || count == 0 || weights == NULL || + grad == NULL) { + return 0; + } + for (size_t i = 0; i < count; i++) { + if (!isfinite(weights[i]) || !isfinite(grad[i])) return 0; + } + for (size_t i = 0; i < count; i++) { + float d = grad[i] + opt->weight_decay * weights[i]; + opt->velocity[i] = opt->momentum * opt->velocity[i] + d; + weights[i] -= opt->lr * opt->velocity[i]; + } + return 1; +} + +void nir_sgd_free(nir_sgd_t *opt) { + if (opt == NULL) return; + free(opt->velocity); + memset(opt, 0, sizeof(*opt)); +} + +int nir_adam_init( + nir_adam_t *opt, size_t count, float lr, float beta1, + float beta2, float eps, float weight_decay, char **error +) { + if (error != NULL) *error = NULL; + if (opt == NULL) { + set_error(error, "optimizer output must not be null"); + return 0; + } + memset(opt, 0, sizeof(*opt)); + if (count == 0 || !isfinite(lr) || lr < 0.0f || + !isfinite(beta1) || beta1 < 0.0f || beta1 >= 1.0f || + !isfinite(beta2) || beta2 < 0.0f || beta2 >= 1.0f || + !isfinite(eps) || eps <= 0.0f || !isfinite(weight_decay) || + weight_decay < 0.0f) { + set_error(error, "invalid Adam hyperparameters"); + return 0; + } + if (count > SIZE_MAX / sizeof(float)) { + set_error(error, "optimizer state is too large"); + return 0; + } + opt->m = (float *)calloc(count, sizeof(float)); + opt->v = (float *)calloc(count, sizeof(float)); + if (opt->m == NULL || opt->v == NULL) { + free(opt->m); + free(opt->v); + memset(opt, 0, sizeof(*opt)); + set_error(error, "out of memory"); + return 0; + } + opt->lr = lr; + opt->beta1 = beta1; + opt->beta2 = beta2; + opt->eps = eps; + opt->weight_decay = weight_decay; + opt->count = count; + opt->step = 0; + return 1; +} + +int nir_adam_step( + nir_adam_t *opt, float *weights, const float *grad, size_t count +) { + if (opt == NULL || opt->m == NULL || opt->v == NULL || + opt->count != count || count == 0 || weights == NULL || + grad == NULL || opt->step == ULONG_MAX) { + return 0; + } + for (size_t i = 0; i < count; i++) { + if (!isfinite(weights[i]) || !isfinite(grad[i])) return 0; + } + opt->step++; + double bc1 = 1.0 - pow((double)opt->beta1, (double)opt->step); + double bc2 = 1.0 - pow((double)opt->beta2, (double)opt->step); + for (size_t i = 0; i < count; i++) { + double d = (double)grad[i] + + (double)opt->weight_decay * (double)weights[i]; + double m = (double)opt->beta1 * (double)opt->m[i] + + (1.0 - (double)opt->beta1) * d; + double v = (double)opt->beta2 * (double)opt->v[i] + + (1.0 - (double)opt->beta2) * d * d; + opt->m[i] = (float)m; + opt->v[i] = (float)v; + double denom = sqrt(v / bc2) + (double)opt->eps; + weights[i] -= (float)((double)opt->lr * (m / bc1) / denom); + } + return 1; +} + +void nir_adam_free(nir_adam_t *opt) { + if (opt == NULL) return; + free(opt->m); + free(opt->v); + memset(opt, 0, sizeof(*opt)); +} + +int nir_stdp_init( + nir_stdp_t *s, size_t in_features, size_t out_features, + float tau_pre, float tau_post, float a_plus, float a_minus, + float w_min, float w_max, char **error +) { + if (error != NULL) *error = NULL; + if (s == NULL) { + set_error(error, "STDP output must not be null"); + return 0; + } + memset(s, 0, sizeof(*s)); + if (in_features == 0 || out_features == 0 || + !isfinite(tau_pre) || tau_pre <= 0.0f || !isfinite(tau_post) || + tau_post <= 0.0f || !isfinite(a_plus) || !isfinite(a_minus) || + !isfinite(w_min) || !isfinite(w_max) || w_min > w_max || + in_features > SIZE_MAX / out_features || + in_features * out_features > SIZE_MAX / sizeof(float)) { + set_error(error, "invalid STDP parameters"); + return 0; + } + s->weight = (float *)calloc( + in_features * out_features, sizeof(float) + ); + s->pre_trace = (float *)calloc(in_features, sizeof(float)); + s->post_trace = (float *)calloc(out_features, sizeof(float)); + if (s->weight == NULL || s->pre_trace == NULL || + s->post_trace == NULL) { + nir_stdp_free(s); + set_error(error, "out of memory"); + return 0; + } + s->in_features = in_features; + s->out_features = out_features; + s->tau_pre = tau_pre; + s->tau_post = tau_post; + s->a_plus = a_plus; + s->a_minus = a_minus; + s->w_min = w_min; + s->w_max = w_max; + return 1; +} + +int nir_stdp_set_weight(nir_stdp_t *s, const float *weight) { + if (s == NULL || s->weight == NULL || weight == NULL) { + return 0; + } + size_t n = s->in_features * s->out_features; + for (size_t i = 0; i < n; i++) { + if (!isfinite(weight[i])) return 0; + } + memcpy(s->weight, weight, n * sizeof(float)); + return 1; +} + +int nir_stdp_step( + nir_stdp_t *s, const float *pre_spikes, const float *post_spikes +) { + if (s == NULL || s->weight == NULL || s->pre_trace == NULL || + s->post_trace == NULL || pre_spikes == NULL || + post_spikes == NULL) { + return 0; + } + size_t in = s->in_features; + size_t out = s->out_features; + for (size_t i = 0; i < in; i++) { + if (!isfinite(pre_spikes[i])) return 0; + } + for (size_t o = 0; o < out; o++) { + if (!isfinite(post_spikes[o])) return 0; + } + /* Decay first: this step's spikes read the prior step's traces. */ + float decay_pre = (float)exp(-1.0 / (double)s->tau_pre); + float decay_post = (float)exp(-1.0 / (double)s->tau_post); + for (size_t i = 0; i < in; i++) s->pre_trace[i] *= decay_pre; + for (size_t o = 0; o < out; o++) s->post_trace[o] *= decay_post; + for (size_t i = 0; i < in; i++) { + for (size_t o = 0; o < out; o++) { + float w = s->weight[i * out + o] + + s->a_plus * s->pre_trace[i] * post_spikes[o] - + s->a_minus * pre_spikes[i] * s->post_trace[o]; + if (w < s->w_min) w = s->w_min; + if (w > s->w_max) w = s->w_max; + s->weight[i * out + o] = w; + } + } + for (size_t i = 0; i < in; i++) s->pre_trace[i] += pre_spikes[i]; + for (size_t o = 0; o < out; o++) s->post_trace[o] += post_spikes[o]; + return 1; +} + +void nir_stdp_free(nir_stdp_t *s) { + if (s == NULL) return; + free(s->weight); + free(s->pre_trace); + free(s->post_trace); + memset(s, 0, sizeof(*s)); +} diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 4e9fbe4..838e77a 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -15,6 +15,7 @@ nir_add_test(test_quantized) nir_add_test(test_workspace) nir_add_test(test_tensor) nir_add_test(test_error_format) +nir_add_test(test_train) # OOM sweep includes src/json.c directly, so it must not link the library # (that would duplicate the json_* symbols); it only needs the headers. diff --git a/tests/test_train.c b/tests/test_train.c new file mode 100644 index 0000000..9c0cc1f --- /dev/null +++ b/tests/test_train.c @@ -0,0 +1,526 @@ +/* C-side training math: surrogate hook, layer backward ops, + * cross-entropy loss, SGD/Adam optimizers, and the STDP port. + * + * Reference values are float64 closed-form evaluations of the + * documented update rules (snnTorch fast_sigmoid/atan gradients, + * torch.optim.SGD/Adam, Song/Abbott STDP, softmax cross-entropy), + * kept to tolerances that separate float32 rounding from formula + * errors. A committed torch-fixture generator under tools/ is + * follow-up work for the tools owner. + */ +#include +#include + +#include "nir_c_runtime/ops.h" +#include "nir_c_runtime/train.h" +#include "test_util.h" + +static nir_tensor_t make_vector(const float *values, size_t n) { + size_t shape[1] = {n}; + nir_tensor_t t = nir_tensor_alloc(shape, 1); + CHECK(t.data != NULL); + memcpy(t.data, values, n * sizeof(float)); + return t; +} + +static void check_vector( + const nir_tensor_t *t, const double *want, size_t n, double tol +) { + CHECK(t->size == n); + for (size_t i = 0; i < n; i++) { + CHECK_NEAR(t->data[i], want[i], tol); + } +} + +static void test_surrogate_names(void) { + nir_surrogate_kind_t kind; + char *error = NULL; + CHECK(nir_surrogate_from_name("fast_sigmoid", &kind, &error)); + CHECK(error == NULL && kind == NIR_SURROGATE_FAST_SIGMOID); + CHECK(nir_surrogate_from_name("atan", &kind, &error)); + CHECK(error == NULL && kind == NIR_SURROGATE_ATAN); + CHECK(!nir_surrogate_from_name("bogus", &kind, &error)); + CHECK(error != NULL && strstr(error, "bogus") != NULL); + free(error); + error = NULL; + CHECK(!nir_surrogate_from_name(NULL, &kind, &error)); + CHECK(error != NULL); + free(error); + CHECK(!nir_surrogate_from_name("atan", NULL, NULL)); +} + +static void test_surrogate_values(void) { + const struct { + nir_surrogate_kind_t kind; + float slope; + float x; + double want; + } cases[] = { + {NIR_SURROGATE_FAST_SIGMOID, 25.0f, -1.0f, 0.03698224852071006}, + {NIR_SURROGATE_FAST_SIGMOID, 25.0f, -0.1f, 2.0408163265306123}, + {NIR_SURROGATE_FAST_SIGMOID, 25.0f, 0.0f, 25.0}, + {NIR_SURROGATE_FAST_SIGMOID, 25.0f, 0.1f, 2.0408163265306123}, + {NIR_SURROGATE_FAST_SIGMOID, 25.0f, 1.0f, 0.03698224852071006}, + {NIR_SURROGATE_ATAN, 2.0f, -1.0f, 0.09199966835037524}, + {NIR_SURROGATE_ATAN, 2.0f, -0.1f, 0.9101698376462755}, + {NIR_SURROGATE_ATAN, 2.0f, 0.0f, 1.0}, + {NIR_SURROGATE_ATAN, 2.0f, 0.1f, 0.9101698376462755}, + {NIR_SURROGATE_ATAN, 2.0f, 1.0f, 0.09199966835037524}, + }; + for (size_t i = 0; i < sizeof(cases) / sizeof(cases[0]); i++) { + float got = 0.0f; + CHECK(nir_surrogate_grad( + cases[i].kind, cases[i].slope, cases[i].x, &got + )); + CHECK_NEAR(got, cases[i].want, 1e-6); + } + float out = 0.0f; + CHECK(!nir_surrogate_grad(NIR_SURROGATE_ATAN, 0.0f, 0.0f, &out)); + CHECK(!nir_surrogate_grad(NIR_SURROGATE_ATAN, -2.0f, 0.0f, &out)); + CHECK(!nir_surrogate_grad( + NIR_SURROGATE_ATAN, (float)INFINITY, 0.0f, &out + )); + CHECK(!nir_surrogate_grad( + NIR_SURROGATE_ATAN, (float)NAN, 0.0f, &out + )); + CHECK(!nir_surrogate_grad( + NIR_SURROGATE_ATAN, 2.0f, (float)NAN, &out + )); + CHECK(!nir_surrogate_grad((nir_surrogate_kind_t)99, 2.0f, 0.0f, &out)); + CHECK(!nir_surrogate_grad(NIR_SURROGATE_ATAN, 2.0f, 0.0f, NULL)); +} + +static void test_threshold_backward(void) { + nir_threshold_params_t params = {0.5f}; + float m[] = {0.2f, 0.5f, 0.8f}; + float g[] = {1.0f, 1.0f, 1.0f}; + nir_tensor_t membrane = make_vector(m, 3); + nir_tensor_t grad_out = make_vector(g, 3); + size_t shape[1] = {3}; + nir_tensor_t grad_in = nir_tensor_alloc(shape, 1); + CHECK(grad_in.data != NULL); + CHECK(nir_op_threshold_backward( + ¶ms, &membrane, &grad_out, NIR_SURROGATE_FAST_SIGMOID, + 25.0f, &grad_in + )); + double want[] = {0.3460207612456747, 25.0, 0.3460207612456747}; + check_vector(&grad_in, want, 3, 1e-6); + /* Exact aliasing follows the elementwise rule. */ + memcpy(grad_out.data, g, sizeof(g)); + CHECK(nir_op_threshold_backward( + ¶ms, &membrane, &grad_out, NIR_SURROGATE_FAST_SIGMOID, + 25.0f, &grad_out + )); + check_vector(&grad_out, want, 3, 1e-6); + CHECK(!nir_op_threshold_backward( + ¶ms, &membrane, &grad_out, NIR_SURROGATE_FAST_SIGMOID, + 0.0f, &grad_in + )); + nir_tensor_free(&membrane); + nir_tensor_free(&grad_out); + nir_tensor_free(&grad_in); +} + +static void test_li_backward(void) { + nir_li_params_t params = {2.0f, 1.0f, 0.0f}; + float g[] = {1.0f, -2.0f}; + nir_tensor_t grad_out = make_vector(g, 2); + size_t shape[1] = {2}; + nir_tensor_t grad_in = nir_tensor_alloc(shape, 1); + nir_tensor_t grad_prev = nir_tensor_alloc(shape, 1); + CHECK(grad_in.data != NULL && grad_prev.data != NULL); + CHECK(nir_op_li_backward(¶ms, &grad_out, &grad_in, &grad_prev)); + double want_in[] = {0.3934693402873666, -0.7869386805747332}; + double want_prev[] = {0.6065306597126334, -1.2130613194252668}; + check_vector(&grad_in, want_in, 2, 1e-6); + check_vector(&grad_prev, want_prev, 2, 1e-6); + CHECK(nir_op_li_backward(¶ms, &grad_out, &grad_in, NULL)); + check_vector(&grad_in, want_in, 2, 1e-6); + nir_li_params_t bad = {0.0f, 1.0f, 0.0f}; + CHECK(!nir_op_li_backward(&bad, &grad_out, &grad_in, NULL)); + nir_tensor_free(&grad_out); + nir_tensor_free(&grad_in); + nir_tensor_free(&grad_prev); +} + +static void test_affine_backward(void) { + float w[] = {1.0f, 2.0f, 3.0f, 4.0f}; + float b[] = {0.5f, -0.5f}; + nir_affine_params_t params = {2, 2, w, b}; + float xv[] = {0.5f, -1.0f}; + float gv[] = {1.0f, -0.5f}; + nir_tensor_t input = make_vector(xv, 2); + nir_tensor_t grad_out = make_vector(gv, 2); + size_t wshape[2] = {2, 2}; + size_t vshape[1] = {2}; + nir_tensor_t grad_w = nir_tensor_alloc(wshape, 2); + nir_tensor_t grad_b = nir_tensor_alloc(vshape, 1); + nir_tensor_t grad_in = nir_tensor_alloc(vshape, 1); + CHECK(grad_w.data != NULL && grad_b.data != NULL); + CHECK(grad_in.data != NULL); + CHECK(nir_op_affine_backward( + ¶ms, &input, &grad_out, &grad_w, &grad_b, &grad_in + )); + double want_w[] = {0.5, -1.0, -0.25, 0.5}; + double want_b[] = {1.0, -0.5}; + double want_x[] = {-0.5, 0.0}; + check_vector(&grad_w, want_w, 4, 1e-6); + check_vector(&grad_b, want_b, 2, 1e-6); + check_vector(&grad_in, want_x, 2, 1e-6); + /* Weight/bias gradients accumulate across calls. */ + CHECK(nir_op_affine_backward( + ¶ms, &input, &grad_out, &grad_w, &grad_b, &grad_in + )); + double want_w2[] = {1.0, -2.0, -0.5, 1.0}; + double want_b2[] = {2.0, -1.0}; + check_vector(&grad_w, want_w2, 4, 1e-6); + check_vector(&grad_b, want_b2, 2, 1e-6); + /* A biased node requires a bias gradient slot. */ + CHECK(!nir_op_affine_backward( + ¶ms, &input, &grad_out, &grad_w, NULL, &grad_in + )); + params.bias = NULL; + memset(grad_w.data, 0, 4 * sizeof(float)); + CHECK(nir_op_affine_backward( + ¶ms, &input, &grad_out, &grad_w, NULL, &grad_in + )); + check_vector(&grad_w, want_w, 4, 1e-6); + nir_tensor_free(&input); + nir_tensor_free(&grad_out); + nir_tensor_free(&grad_w); + nir_tensor_free(&grad_b); + nir_tensor_free(&grad_in); +} + +static void test_linear_backward(void) { + float w[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; + nir_linear_params_t params = {2, 3, w}; + float xv[] = {1.0f, 0.0f, -1.0f}; + float gv[] = {2.0f, -1.0f}; + nir_tensor_t input = make_vector(xv, 3); + nir_tensor_t grad_out = make_vector(gv, 2); + size_t wshape[2] = {2, 3}; + size_t xshape[1] = {3}; + nir_tensor_t grad_w = nir_tensor_alloc(wshape, 2); + nir_tensor_t grad_in = nir_tensor_alloc(xshape, 1); + CHECK(grad_w.data != NULL && grad_in.data != NULL); + CHECK(nir_op_linear_backward( + ¶ms, &input, &grad_out, &grad_w, &grad_in + )); + double want_w[] = {2.0, 0.0, -2.0, -1.0, 0.0, 1.0}; + double want_x[] = {-2.0, -1.0, 0.0}; + check_vector(&grad_w, want_w, 6, 1e-6); + check_vector(&grad_in, want_x, 3, 1e-6); + nir_tensor_free(&input); + nir_tensor_free(&grad_out); + nir_tensor_free(&grad_w); + nir_tensor_free(&grad_in); +} + +static void test_cross_entropy(void) { + float lv[] = {1.0f, 2.0f, 3.0f}; + nir_tensor_t logits = make_vector(lv, 3); + size_t shape[1] = {3}; + nir_tensor_t grad = nir_tensor_alloc(shape, 1); + CHECK(grad.data != NULL); + float loss = 0.0f; + CHECK(nir_cross_entropy(&logits, 2, &loss, &grad)); + CHECK_NEAR(loss, 0.40760596444438046, 1e-6); + double want[] = {0.09003057317038046, 0.24472847105479764, + -0.3347590442251782}; + check_vector(&grad, want, 3, 1e-6); + CHECK(!nir_cross_entropy(&logits, 3, &loss, &grad)); + nir_tensor_free(&logits); + nir_tensor_free(&grad); +} + +static void test_sgd(void) { + nir_sgd_t opt; + char *error = NULL; + CHECK(nir_sgd_init(&opt, 2, 0.1f, 0.9f, 0.01f, &error)); + CHECK(error == NULL); + float w[] = {1.0f, -2.0f}; + float g[] = {0.1f, 0.2f}; + CHECK(nir_sgd_step(&opt, w, g, 2)); + CHECK_NEAR(w[0], 0.989, 1e-6); + CHECK_NEAR(w[1], -2.018, 1e-6); + CHECK(nir_sgd_step(&opt, w, g, 2)); + CHECK_NEAR(w[0], 0.9681109999999999, 1e-6); + CHECK_NEAR(w[1], -2.0521819999999997, 1e-6); + CHECK(!nir_sgd_step(&opt, w, g, 1)); + nir_sgd_free(&opt); + CHECK(nir_sgd_init(&opt, 1, 0.5f, 0.0f, 0.0f, &error)); + float w0[] = {2.0f}; + float g0[] = {1.0f}; + CHECK(nir_sgd_step(&opt, w0, g0, 1)); + CHECK_NEAR(w0[0], 1.5, 1e-6); + nir_sgd_free(&opt); + CHECK(!nir_sgd_init(&opt, 2, -0.1f, 0.0f, 0.0f, &error)); + CHECK(error != NULL); + free(error); +} + +static void test_adam(void) { + nir_adam_t opt; + char *error = NULL; + CHECK(nir_adam_init( + &opt, 2, 0.001f, 0.9f, 0.999f, 1e-8f, 0.0f, &error + )); + CHECK(error == NULL); + float w[] = {1.0f, -2.0f}; + float g[] = {0.1f, 0.2f}; + CHECK(nir_adam_step(&opt, w, g, 2)); + CHECK_NEAR(w[0], 0.9990000001, 1e-6); + CHECK_NEAR(w[1], -2.00099999995, 1e-6); + CHECK(nir_adam_step(&opt, w, g, 2)); + CHECK_NEAR(w[0], 0.9980000002, 1e-6); + CHECK_NEAR(w[1], -2.0019999998999998, 1e-6); + nir_adam_free(&opt); + CHECK(nir_adam_init( + &opt, 2, 0.001f, 0.9f, 0.999f, 1e-8f, 0.01f, &error + )); + float w2[] = {1.0f, -2.0f}; + CHECK(nir_adam_step(&opt, w2, g, 2)); + CHECK_NEAR(w2[0], 0.9990000000909091, 1e-6); + CHECK_NEAR(w2[1], -2.0009999999444443, 1e-6); + nir_adam_free(&opt); + CHECK(!nir_adam_init( + &opt, 2, 0.001f, 1.0f, 0.999f, 1e-8f, 0.0f, &error + )); + CHECK(error != NULL); + free(error); +} + +static void test_stdp(void) { + nir_stdp_t s; + char *error = NULL; + CHECK(nir_stdp_init( + &s, 2, 2, 2.0f, 2.0f, 0.01f, 0.012f, 0.0f, 1.0f, &error + )); + CHECK(error == NULL); + float pre1[] = {1.0f, 0.0f}; + float post1[] = {0.0f, 1.0f}; + CHECK(nir_stdp_step(&s, pre1, post1)); + CHECK_NEAR(s.weight[0], 0.0, 1e-7); + CHECK_NEAR(s.weight[3], 0.0, 1e-7); + float pre2[] = {0.0f, 1.0f}; + float post2[] = {1.0f, 0.0f}; + CHECK(nir_stdp_step(&s, pre2, post2)); + /* w[0][0] potentiates; w[1][1]'s raw -0.00727... clamps to 0. */ + CHECK_NEAR(s.weight[0], 0.006065306597126334, 1e-7); + CHECK_NEAR(s.weight[1], 0.0, 1e-7); + CHECK_NEAR(s.weight[2], 0.0, 1e-7); + CHECK_NEAR(s.weight[3], 0.0, 1e-7); + CHECK_NEAR(s.pre_trace[0], 0.6065306597126334, 1e-6); + CHECK_NEAR(s.pre_trace[1], 1.0, 1e-6); + nir_stdp_free(&s); + /* Upper clamp: the raw update overshoots w_max. */ + CHECK(nir_stdp_init( + &s, 1, 1, 10.0f, 10.0f, 0.1f, 0.012f, 0.0f, 0.05f, &error + )); + float one[] = {1.0f}; + float zero[] = {0.0f}; + CHECK(nir_stdp_step(&s, one, zero)); + CHECK(nir_stdp_step(&s, zero, one)); + CHECK_NEAR(s.weight[0], 0.05, 1e-6); + float seed[] = {0.02f}; + CHECK(nir_stdp_set_weight(&s, seed)); + CHECK_NEAR(s.weight[0], 0.02, 1e-7); + nir_stdp_free(&s); + CHECK(!nir_stdp_init( + &s, 1, 1, 0.0f, 10.0f, 0.1f, 0.012f, 0.0f, 1.0f, &error + )); + CHECK(error != NULL); + free(error); +} + +/* A two-layer topology: Affine(2->2) + Threshold(0.5) + + * Affine(2->2), trained single-step against one sample. */ +typedef struct { + float w1[4]; + float b1[2]; + float w2[4]; + float b2[2]; + nir_affine_params_t l1; + nir_affine_params_t l2; + nir_threshold_params_t thr; +} topo_t; + +static void topo_init(topo_t *net) { + float w1[] = {0.5f, -0.5f, 0.25f, 0.75f}; + float b1[] = {0.0f, 0.1f}; + float w2[] = {1.0f, -1.0f, -0.5f, 0.5f}; + float b2[] = {0.1f, -0.1f}; + memcpy(net->w1, w1, sizeof(w1)); + memcpy(net->b1, b1, sizeof(b1)); + memcpy(net->w2, w2, sizeof(w2)); + memcpy(net->b2, b2, sizeof(b2)); + net->l1.out_features = 2; + net->l1.in_features = 2; + net->l1.weight = net->w1; + net->l1.bias = net->b1; + net->l2.out_features = 2; + net->l2.in_features = 2; + net->l2.weight = net->w2; + net->l2.bias = net->b2; + net->thr.threshold = 0.5f; +} + +typedef struct { + nir_tensor_t h_pre; + nir_tensor_t h; + nir_tensor_t logits; + nir_tensor_t dlogits; + nir_tensor_t dh; + nir_tensor_t dh_pre; + nir_tensor_t dw1; + nir_tensor_t db1; + nir_tensor_t dw2; + nir_tensor_t db2; +} topo_grad_t; + +static void topo_alloc(topo_grad_t *g) { + size_t v[1] = {2}; + size_t m[2] = {2, 2}; + g->h_pre = nir_tensor_alloc(v, 1); + g->h = nir_tensor_alloc(v, 1); + g->logits = nir_tensor_alloc(v, 1); + g->dlogits = nir_tensor_alloc(v, 1); + g->dh = nir_tensor_alloc(v, 1); + g->dh_pre = nir_tensor_alloc(v, 1); + g->dw1 = nir_tensor_alloc(m, 2); + g->db1 = nir_tensor_alloc(v, 1); + g->dw2 = nir_tensor_alloc(m, 2); + g->db2 = nir_tensor_alloc(v, 1); +} + +static void topo_free(topo_grad_t *g) { + nir_tensor_free(&g->h_pre); + nir_tensor_free(&g->h); + nir_tensor_free(&g->logits); + nir_tensor_free(&g->dlogits); + nir_tensor_free(&g->dh); + nir_tensor_free(&g->dh_pre); + nir_tensor_free(&g->dw1); + nir_tensor_free(&g->db1); + nir_tensor_free(&g->dw2); + nir_tensor_free(&g->db2); +} + +static float topo_loss_and_grads( + topo_t *net, topo_grad_t *g, const nir_tensor_t *input, + size_t target +) { + memset(g->dw1.data, 0, 4 * sizeof(float)); + memset(g->db1.data, 0, 2 * sizeof(float)); + memset(g->dw2.data, 0, 4 * sizeof(float)); + memset(g->db2.data, 0, 2 * sizeof(float)); + CHECK(nir_op_affine_into(&net->l1, input, &g->h_pre)); + CHECK(nir_op_threshold_into(&net->thr, &g->h_pre, &g->h)); + CHECK(nir_op_affine_into(&net->l2, &g->h, &g->logits)); + float loss = 0.0f; + CHECK(nir_cross_entropy(&g->logits, target, &loss, &g->dlogits)); + CHECK(nir_op_affine_backward( + &net->l2, &g->h, &g->dlogits, &g->dw2, &g->db2, &g->dh + )); + CHECK(nir_op_threshold_backward( + &net->thr, &g->h_pre, &g->dh, NIR_SURROGATE_FAST_SIGMOID, + 25.0f, &g->dh_pre + )); + size_t v[1] = {2}; + nir_tensor_t dx = nir_tensor_alloc(v, 1); + CHECK(dx.data != NULL); + CHECK(nir_op_affine_backward( + &net->l1, input, &g->dh_pre, &g->dw1, &g->db1, &dx + )); + nir_tensor_free(&dx); + return loss; +} + +static void test_topology_step(void) { + topo_t net; + topo_init(&net); + topo_grad_t g; + topo_alloc(&g); + float xv[] = {0.5f, -0.25f}; + nir_tensor_t input = make_vector(xv, 2); + float loss = topo_loss_and_grads(&net, &g, &input, 1); + CHECK_NEAR(loss, 0.7981388693815918, 1e-5); + char *error = NULL; + nir_adam_t o_w1, o_b1, o_w2, o_b2; + CHECK(nir_adam_init(&o_w1, 4, 0.01f, 0.9f, 0.999f, 1e-8f, 0, &error)); + CHECK(nir_adam_init(&o_b1, 2, 0.01f, 0.9f, 0.999f, 1e-8f, 0, &error)); + CHECK(nir_adam_init(&o_w2, 4, 0.01f, 0.9f, 0.999f, 1e-8f, 0, &error)); + CHECK(nir_adam_init(&o_b2, 2, 0.01f, 0.9f, 0.999f, 1e-8f, 0, &error)); + CHECK(nir_adam_step(&o_w1, net.w1, g.dw1.data, 4)); + CHECK(nir_adam_step(&o_b1, net.b1, g.db1.data, 2)); + CHECK(nir_adam_step(&o_w2, net.w2, g.dw2.data, 4)); + CHECK(nir_adam_step(&o_b2, net.b2, g.db2.data, 2)); + double want[] = {0.4900000001650498, -0.4900000003300996, + 0.2599999984691973, 0.740000003061605, -0.009999999917475093, + 0.10999999923459859, 1.0, -1.0, -0.5, 0.5, + 0.09000000018187308, -0.09000000018187308}; + const float *bufs[] = {net.w1, net.b1, net.w2, net.b2}; + size_t sizes[] = {4, 2, 4, 2}; + size_t k = 0; + for (size_t b = 0; b < 4; b++) { + for (size_t i = 0; i < sizes[b]; i++, k++) { + CHECK_NEAR(bufs[b][i], want[k], 1e-5); + } + } + loss = topo_loss_and_grads(&net, &g, &input, 1); + CHECK_NEAR(loss, 0.7871917250389795, 1e-5); + nir_adam_free(&o_w1); + nir_adam_free(&o_b1); + nir_adam_free(&o_w2); + nir_adam_free(&o_b2); + nir_tensor_free(&input); + topo_free(&g); +} + +static void test_topology_converges(void) { + topo_t net; + topo_init(&net); + topo_grad_t g; + topo_alloc(&g); + float xv[] = {0.5f, -0.25f}; + nir_tensor_t input = make_vector(xv, 2); + char *error = NULL; + nir_adam_t o_w1, o_b1, o_w2, o_b2; + CHECK(nir_adam_init(&o_w1, 4, 0.05f, 0.9f, 0.999f, 1e-8f, 0, &error)); + CHECK(nir_adam_init(&o_b1, 2, 0.05f, 0.9f, 0.999f, 1e-8f, 0, &error)); + CHECK(nir_adam_init(&o_w2, 4, 0.05f, 0.9f, 0.999f, 1e-8f, 0, &error)); + CHECK(nir_adam_init(&o_b2, 2, 0.05f, 0.9f, 0.999f, 1e-8f, 0, &error)); + float loss = 0.0f; + for (int step = 0; step < 300; step++) { + loss = topo_loss_and_grads(&net, &g, &input, 1); + CHECK(nir_adam_step(&o_w1, net.w1, g.dw1.data, 4)); + CHECK(nir_adam_step(&o_b1, net.b1, g.db1.data, 2)); + CHECK(nir_adam_step(&o_w2, net.w2, g.dw2.data, 4)); + CHECK(nir_adam_step(&o_b2, net.b2, g.db2.data, 2)); + } + loss = topo_loss_and_grads(&net, &g, &input, 1); + CHECK(loss < 0.01f); + nir_adam_free(&o_w1); + nir_adam_free(&o_b1); + nir_adam_free(&o_w2); + nir_adam_free(&o_b2); + nir_tensor_free(&input); + topo_free(&g); +} + +int main(void) { + test_surrogate_names(); + test_surrogate_values(); + test_threshold_backward(); + test_li_backward(); + test_affine_backward(); + test_linear_backward(); + test_cross_entropy(); + test_sgd(); + test_adam(); + test_stdp(); + test_topology_step(); + test_topology_converges(); + return 0; +}