From 854998d6238bea74d0c2ff0b571232cf6ab99400 Mon Sep 17 00:00:00 2001 From: w4ffl35 <25737761+w4ffl35@users.noreply.github.com> Date: Wed, 14 Oct 2026 07:06:17 -0600 Subject: [PATCH] Add C training driver plus worked MNIST example (closes #44) driver.h/driver.c/driver_loop.c: nir_driver_t wires IDX loading, repeat/rate encoding, tape forward/backward, minibatch mean reduction, SGD/Adam steps with LR schedules, and checkpoint save/load with optimizer resume. Per-node accumulators and optimizer instances, seeded shuffle/encode streams, step counter persisted as checkpoint meta step. examples/mnist_train plus mnist_tiny_graph.json: 4-step LI network (Flatten 4->Affine 8->LI 8->Affine 10) trains the committed 2x2 3-sample fixture for 20 Adam epochs, printing loss/accuracy per epoch and saving a v2 checkpoint. Deterministic from seed 7: 2.39/0.000 to 1.16/0.667, step 40. Tests: test_driver asserts 20-epoch loss beats init by 0.5, checkpoint resume is bit-identical with matching step, SGD/rate/schedule smoke paths run finite, and config/data error paths fail cleanly. mnist_train_example runs the binary via ctest; CI also runs it with PATH=/nonexistent to prove zero Python at train time. Docs: README Training section with flow diagram, example invocation, expected output, and zero-Python proof; Scope marks tape/backward/minibatch/checkpoint/init/schedule/loss/driver done. train.h comment refresh for tape/driver. --- .github/workflows/ci.yml | 7 + CMakeLists.txt | 3 + README.md | 81 ++++- examples/CMakeLists.txt | 12 + examples/mnist_tiny_graph.json | 241 ++++++++++++ examples/mnist_train.c | 223 ++++++++++++ include/nir_c_runtime/driver.h | 156 ++++++++ include/nir_c_runtime/train.h | 11 +- src/driver.c | 324 +++++++++++++++++ src/driver_loop.c | 646 +++++++++++++++++++++++++++++++++ tests/CMakeLists.txt | 1 + tests/test_driver.c | 303 ++++++++++++++++ 12 files changed, 1983 insertions(+), 25 deletions(-) create mode 100644 examples/CMakeLists.txt create mode 100644 examples/mnist_tiny_graph.json create mode 100644 examples/mnist_train.c create mode 100644 include/nir_c_runtime/driver.h create mode 100644 src/driver.c create mode 100644 src/driver_loop.c create mode 100644 tests/test_driver.c diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2586380..d975a93 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -29,6 +29,13 @@ jobs: run: cmake --build build --parallel - name: Test run: ctest --test-dir build --output-on-failure + - name: Train with no Python on PATH + run: >- + PATH=/nonexistent ./build/examples/mnist_train + examples/mnist_tiny_graph.json + tests/fixtures/mnist_tiny_images.idx + tests/fixtures/mnist_tiny_labels.idx + ./build/mnist_nopython_checkpoint.json 2 sanitizers: runs-on: ubuntu-latest diff --git a/CMakeLists.txt b/CMakeLists.txt index cabae61..4c9187c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -28,9 +28,12 @@ add_library(nir_c_runtime src/checkpoint.c src/minibatch.c src/tape.c + src/driver.c + src/driver_loop.c ) target_include_directories(nir_c_runtime PUBLIC include) target_link_libraries(nir_c_runtime PUBLIC m) enable_testing() add_subdirectory(tests) +add_subdirectory(examples) diff --git a/README.md b/README.md index 75e8d4f..65d3aa7 100644 --- a/README.md +++ b/README.md @@ -84,29 +84,70 @@ target-specific kernels remain future optimization work. spike encoders (`encode.h`), IDX MNIST dataset loading (`data.h`), and versioned checkpoint save/load with weight accessors and summary export (`checkpoint.h`), all covered by C tests. +- [x] Backward ops for spatial/linear nodes (Conv2d, Conv1d, Flatten, + AvgPool2d, SumPool2d, Scale) and spiking neurons (LIF/IF/CubaLIF + with surrogate through the spike output, zero through the + hard-reset branch). +- [x] Graph-level BPTT tape (`tape.h`): retains per-step activations + and neuron state, backpropagates through Delay/LI recurrence, + matches `nir_runtime_run` bit for bit, finite-difference + verified on smooth topologies. +- [x] Mini-batch training (`minibatch.h`): mean loss/gradient + reduction over a batch-1 loop (no batch axis on the tensor + ABI) plus seeded Fisher-Yates epoch shuffling. +- [x] Optimizer-state checkpoints: v2 `optimizers` section persists + SGD velocity and Adam m/v/step per node for bit-identical + resume; v1 files still load. +- [x] Weight initializers (`init.h`: Xavier/He/uniform from + `nir_rng_t`), LR schedules (`schedule.h`: step/cosine/ + exponential mirroring torch), and extra losses (`loss.h`: + MSE plus spike-count cross-entropy). +- [x] End-to-end C training driver (`driver.h`) plus a worked MNIST + example (`examples/mnist_train`): loads IDX, encodes, runs + the tape, steps SGD/Adam per batch, saves optimizer + checkpoints — zero Python at train time (see Training). + +## Training + +The driver wires the tested units into one loop; no Python or +torch is involved after the fixtures are committed: + +```text +IDX MNIST (data.h) → encode (encode.h: repeat/rate) + → tape forward (tape.h, unprimed per sample) + → cross-entropy (train.h) → tape backward + → minibatch mean (minibatch.h) → SGD/Adam step + → checkpoint with optimizer state (checkpoint.h) +``` + +Build and run the worked example (tiny 2×2, 3-sample fixture, +4-step LI network, 20 Adam epochs): + +```sh +cmake -S . -B build +cmake --build build --parallel +./build/examples/mnist_train \ + examples/mnist_tiny_graph.json \ + tests/fixtures/mnist_tiny_images.idx \ + tests/fixtures/mnist_tiny_labels.idx \ + /tmp/mnist_tiny_checkpoint.json 20 +``` + +Expected output (deterministic from seed 7): initial loss +~2.39 at 0.000 accuracy, falling to ~1.16 at 0.667 accuracy +by epoch 20, saving a v2 checkpoint with `step 40` (2 batches +× 20 epochs). `tests/test_driver.c` asserts the same run beats +its initial loss by 0.5 and resumes bit-identically from the +saved checkpoint. + +Zero-Python proof: the example is a pure C binary — CI runs it +via ctest and again with `PATH=/nonexistent` (no `python` or +`python3` resolvable), and it still trains and saves. The only +Python in the repo (`tools/`) regenerates fixtures; training +and inference never invoke it. ## Out of scope (for now) -- Full training pipelines. The training math and training IO above are - supported; still missing are the graph-level backprop-through-time tape - ([#39](https://github.com/Capsize-Games/nir-c-runtime/issues/39)), - backward ops for spatial/linear - ([#40](https://github.com/Capsize-Games/nir-c-runtime/issues/40)) and - spiking-neuron - ([#41](https://github.com/Capsize-Games/nir-c-runtime/issues/41)) nodes, - mini-batch training - ([#42](https://github.com/Capsize-Games/nir-c-runtime/issues/42)), - optimizer-state checkpoint resume - ([#43](https://github.com/Capsize-Games/nir-c-runtime/issues/43)), and an - end-to-end C training driver with a worked example - ([#44](https://github.com/Capsize-Games/nir-c-runtime/issues/44)). - Weight initializers - ([#45](https://github.com/Capsize-Games/nir-c-runtime/issues/45)), - learning-rate schedules - ([#46](https://github.com/Capsize-Games/nir-c-runtime/issues/46)), and - additional loss functions - ([#47](https://github.com/Capsize-Games/nir-c-runtime/issues/47)) are filed - as follow-ups. - 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/examples/CMakeLists.txt b/examples/CMakeLists.txt new file mode 100644 index 0000000..1ee28d0 --- /dev/null +++ b/examples/CMakeLists.txt @@ -0,0 +1,12 @@ +add_executable(mnist_train mnist_train.c) +target_link_libraries(mnist_train PRIVATE nir_c_runtime) + +add_test( + NAME mnist_train_example + COMMAND mnist_train + ${CMAKE_SOURCE_DIR}/examples/mnist_tiny_graph.json + ${CMAKE_SOURCE_DIR}/tests/fixtures/mnist_tiny_images.idx + ${CMAKE_SOURCE_DIR}/tests/fixtures/mnist_tiny_labels.idx + ${CMAKE_BINARY_DIR}/mnist_tiny_checkpoint.json + 20 +) diff --git a/examples/mnist_tiny_graph.json b/examples/mnist_tiny_graph.json new file mode 100644 index 0000000..a0a96a4 --- /dev/null +++ b/examples/mnist_tiny_graph.json @@ -0,0 +1,241 @@ +{ + "nodes": [ + { + "name": "input", + "kind": "Input", + "params": {} + }, + { + "name": "flatten", + "kind": "Flatten", + "params": {} + }, + { + "name": "fc1", + "kind": "Affine", + "params": { + "weight": [ + [ + -0.24913809321871228, + -0.4937738335497391, + 0.2134535788025974, + -0.6046664021733632 + ], + [ + 0.05074481713564971, + -0.18994455527924115, + -0.6250839151671003, + 0.010515714722665881 + ], + [ + -0.6540799124877833, + -0.09383917408664855, + -0.6083162937620042, + -0.5788192074319221 + ], + [ + -0.10674598641357513, + 0.46223870760165797, + -0.5320243686803318, + -0.39139920978918025 + ], + [ + 0.1802177914228914, + 0.6331560584184129, + 0.10904003563382247, + -0.14611587400681747 + ], + [ + 0.6735264294789381, + -0.6412289224852223, + 0.506950956469627, + -0.29753740066709794 + ], + [ + -0.503099285861198, + -0.5405234005538385, + -0.2708476017960004, + 0.4470701844911761 + ], + [ + -0.45152108361952253, + 0.11540005814332399, + 0.196453311751703, + -0.1804571256694043 + ] + ], + "bias": [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ] + } + }, + { + "name": "li1", + "kind": "LI", + "params": { + "tau": 2.0, + "r": 1.0, + "v_leak": 0.0 + } + }, + { + "name": "fc2", + "kind": "Affine", + "params": { + "weight": [ + [ + 0.055130560259456174, + -0.5048478059836482, + -0.5085287661415836, + -0.33952963261325175, + 0.20830794615661075, + -0.08360920362634078, + -0.21460436242522613, + 0.09879832985700243 + ], + [ + -0.05405802580922614, + -0.23120915652280033, + 0.3399201458018255, + 0.22977897976201023, + -0.29549189684225907, + 0.08593709830386886, + 0.0290944165163578, + 0.4331714681048824 + ], + [ + 0.2649405992440478, + -0.244868377051209, + 0.5544581549158596, + -0.44101965147445565, + -0.09454362176557207, + 0.2969207698083346, + -0.4018536451918209, + -0.012744313822624975 + ], + [ + -0.5320776283685758, + 0.19423894010419307, + 0.3055001216554002, + 0.08432309255394632, + 0.43356443147059, + -0.2150658471887773, + 0.2255076645791675, + 0.1089689478999486 + ], + [ + 0.092255035398914, + -0.050569727524399255, + 0.3925609791894308, + 0.5134734999282105, + -0.029908663726454288, + 0.18954664003774813, + -0.5072952484799551, + 0.23266294547934185 + ], + [ + 0.1698897675342217, + 0.5693781467747452, + 0.3717267244158634, + -0.24872765506018207, + -0.13187668289430077, + 0.19474338183059503, + -0.5512968440164253, + -0.044230473532123504 + ], + [ + -0.38330471559247825, + -0.4421396922596237, + -0.5092755694479254, + 0.30972877620029526, + -0.42800104519055915, + -0.29142928740911, + -0.12592043650247992, + 0.42888115348952227 + ], + [ + -0.48430299731052506, + -0.05867333548025655, + 0.057088289706041295, + 0.44269351079791197, + 0.36867260064257057, + 0.42029306312256176, + -0.2558573160993044, + -0.09780715717824934 + ], + [ + -0.16307701142624714, + 0.4436276644072279, + 0.5285424676502037, + -0.4030818180199479, + -0.3738715632298264, + -0.30950955019234344, + -0.3079169677400293, + -0.01736354337058954 + ], + [ + 0.10291095774188086, + -0.27395660642827574, + -0.5726233831569812, + -0.09359251878810587, + -0.15097296976962415, + 0.07660424673050914, + 0.5231924185423495, + 0.21996312845274624 + ] + ], + "bias": [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ] + } + }, + { + "name": "output", + "kind": "Output", + "params": {} + } + ], + "edges": [ + { + "source": "input", + "target": "flatten", + "delayed": false + }, + { + "source": "flatten", + "target": "fc1", + "delayed": false + }, + { + "source": "fc1", + "target": "li1", + "delayed": false + }, + { + "source": "li1", + "target": "fc2", + "delayed": false + }, + { + "source": "fc2", + "target": "output", + "delayed": false + } + ] +} diff --git a/examples/mnist_train.c b/examples/mnist_train.c new file mode 100644 index 0000000..47e277d --- /dev/null +++ b/examples/mnist_train.c @@ -0,0 +1,223 @@ +/* Worked MNIST training example (#44): zero Python at + * train time. Loads a committed graph plus the tiny IDX + * fixture, trains with the C driver (encode, tape BPTT, + * minibatch mean, Adam), prints loss/accuracy per epoch, + * and saves a checkpoint with optimizer state. + * + * Usage: mnist_train + * [epochs] + */ +#include +#include +#include +#include + +#include "nir_c_runtime/checkpoint.h" +#include "nir_c_runtime/data.h" +#include "nir_c_runtime/driver.h" +#include "nir_c_runtime/graph.h" +#include "nir_c_runtime/json.h" +#include "nir_c_runtime/train.h" + +static char *read_whole(const char *path) { + FILE *f = fopen(path, "rb"); + if (f == NULL) { + return NULL; + } + fseek(f, 0, SEEK_END); + long size = ftell(f); + fseek(f, 0, SEEK_SET); + if (size < 0) { + fclose(f); + return NULL; + } + char *buf = (char *)malloc((size_t)size + 1); + if (buf == NULL) { + fclose(f); + return NULL; + } + if (size > 0 && + fread(buf, 1, (size_t)size, f) != (size_t)size) { + free(buf); + fclose(f); + return NULL; + } + buf[size] = '\0'; + fclose(f); + return buf; +} + +static int load_graph( + const char *path, nir_graph_t *graph +) { + char *text = read_whole(path); + if (text == NULL) { + fprintf(stderr, "cannot open graph: %s\n", path); + return 0; + } + char *error = NULL; + json_value_t *json = json_parse(text, &error); + if (error != NULL) { + fprintf(stderr, "graph parse failed: %s\n", error); + free(error); + free(text); + return 0; + } + memset(graph, 0, sizeof(*graph)); + int ok = nir_graph_load(json, graph, &error); + if (!ok) { + fprintf(stderr, "graph load failed: %s\n", + error != NULL ? error : "unknown"); + free(error); + } + json_free(json); + free(text); + return ok; +} + +static void fill_config(nir_driver_config_t *c) { + memset(c, 0, sizeof(*c)); + c->num_steps = 4; + c->surrogate = NIR_SURROGATE_FAST_SIGMOID; + c->slope = 25.0f; + c->encode = NIR_DRIVER_ENCODE_REPEAT; + c->encode_gain = 1.0f; + c->batch_size = 2; + c->seed = 7; + c->opt_kind = NIR_DRIVER_OPT_ADAM; + c->lr = 0.01f; + c->momentum = 0.0f; + c->weight_decay = 0.0f; + c->beta1 = 0.9f; + c->beta2 = 0.999f; + c->eps = 1e-8f; + c->sched = NIR_DRIVER_SCHED_NONE; +} + +static int load_dataset( + const char *img_path, const char *lbl_path, + uint8_t **px, uint8_t **lb, size_t *count, + size_t *rows, size_t *cols +) { + char *error = NULL; + size_t nlabels = 0; + if (!nir_idx_read_images(img_path, px, count, rows, + cols, &error)) { + fprintf(stderr, "images failed: %s\n", + error != NULL ? error : "unknown"); + free(error); + return 0; + } + if (!nir_idx_read_labels(lbl_path, lb, &nlabels, + &error)) { + fprintf(stderr, "labels failed: %s\n", + error != NULL ? error : "unknown"); + free(error); + free(*px); + *px = NULL; + return 0; + } + if (nlabels != *count) { + fprintf(stderr, "images/labels count mismatch\n"); + free(*px); + free(*lb); + *px = NULL; + *lb = NULL; + return 0; + } + return 1; +} + +int main(int argc, char **argv) { + if (argc < 5 || argc > 6) { + fprintf(stderr, + "usage: %s " + " [epochs]\n", + argv[0]); + return 2; + } + int epochs = 20; + if (argc == 6) { + epochs = atoi(argv[5]); + if (epochs <= 0 || epochs > 10000) { + fprintf(stderr, "epochs must be 1..10000\n"); + return 2; + } + } + nir_graph_t graph; + if (!load_graph(argv[1], &graph)) { + return 1; + } + uint8_t *px = NULL; + uint8_t *lb = NULL; + size_t count = 0; + size_t rows = 0; + size_t cols = 0; + if (!load_dataset(argv[2], argv[3], &px, &lb, &count, + &rows, &cols)) { + nir_graph_free(&graph); + return 1; + } + nir_driver_config_t config; + fill_config(&config); + nir_driver_t driver; + memset(&driver, 0, sizeof(driver)); + char *error = NULL; + if (!nir_driver_init(&driver, &graph, &config, + &error)) { + fprintf(stderr, "driver init failed: %s\n", + error != NULL ? error : "unknown"); + free(error); + free(px); + free(lb); + nir_graph_free(&graph); + return 1; + } + nir_driver_stats_t stats; + if (!nir_driver_eval(&driver, px, lb, count, rows, + cols, &stats, &error)) { + fprintf(stderr, "initial eval failed: %s\n", + error != NULL ? error : "unknown"); + free(error); + nir_driver_free(&driver); + free(px); + free(lb); + nir_graph_free(&graph); + return 1; + } + printf("initial loss %.6f acc %.3f\n", stats.loss, + stats.accuracy); + for (int e = 0; e < epochs; e++) { + if (!nir_driver_train_epoch(&driver, px, lb, + count, rows, cols, &stats, &error)) { + fprintf(stderr, "epoch %d failed: %s\n", e, + error != NULL ? error : "unknown"); + free(error); + nir_driver_free(&driver); + free(px); + free(lb); + nir_graph_free(&graph); + return 1; + } + printf("epoch %d loss %.6f acc %.3f\n", e + 1, + stats.loss, stats.accuracy); + } + if (!nir_driver_save(&driver, argv[4], "mnist_tiny", + &error)) { + fprintf(stderr, "checkpoint save failed: %s\n", + error != NULL ? error : "unknown"); + free(error); + nir_driver_free(&driver); + free(px); + free(lb); + nir_graph_free(&graph); + return 1; + } + printf("saved %s step %lu\n", argv[4], + nir_driver_step(&driver)); + nir_driver_free(&driver); + free(px); + free(lb); + nir_graph_free(&graph); + return 0; +} diff --git a/include/nir_c_runtime/driver.h b/include/nir_c_runtime/driver.h new file mode 100644 index 0000000..c06d94d --- /dev/null +++ b/include/nir_c_runtime/driver.h @@ -0,0 +1,156 @@ +/* End-to-end C training driver (#44). + * + * Wires the tested training pieces into one loop with zero + * Python at train time: IDX loading (data.h), spike encoding + * (encode.h), BPTT tape forward/backward (tape.h), minibatch + * mean reduction (minibatch.h), SGD/Adam steps (train.h), LR + * schedules (schedule.h), and checkpoint save/load with + * optimizer resume (checkpoint.h). + * + * The graph is borrowed and mutated in place (optimizer steps + * write its weight/bias buffers). The tape owns retained + * activations; the driver owns per-node accumulators and + * optimizer instances. Temporal state never carries across + * samples: every sample pairs one fresh tape forward with one + * backward, matching the tape's unprimed convention. + * + * Batch reduction follows the minibatch.h contract: per-sample + * gradients accumulate, scale by 1/B, then step. The LR + * schedule counter is the driver's completed-step count, + * persisted as the checkpoint meta step so resume continues + * the curve exactly. + */ +#ifndef NIR_C_RUNTIME_DRIVER_H +#define NIR_C_RUNTIME_DRIVER_H + +#include +#include + +#include "nir_c_runtime/ops.h" +#include "nir_c_runtime/rng.h" +#include "nir_c_runtime/tape.h" +#include "nir_c_runtime/train.h" + +typedef enum { + NIR_DRIVER_OPT_SGD, + NIR_DRIVER_OPT_ADAM +} nir_driver_opt_kind_t; + +typedef enum { + NIR_DRIVER_ENCODE_REPEAT, + NIR_DRIVER_ENCODE_RATE +} nir_driver_encode_t; + +typedef enum { + NIR_DRIVER_SCHED_NONE, + NIR_DRIVER_SCHED_STEP, + NIR_DRIVER_SCHED_COSINE, + NIR_DRIVER_SCHED_EXP +} nir_driver_sched_t; + +typedef struct { + size_t num_steps; + nir_surrogate_kind_t surrogate; + float slope; + nir_driver_encode_t encode; + float encode_gain; + size_t batch_size; + uint64_t seed; + nir_driver_opt_kind_t opt_kind; + float lr; + float momentum; + float weight_decay; + float beta1; + float beta2; + float eps; + nir_driver_sched_t sched; + float sched_gamma; + unsigned long sched_step_size; + float sched_eta_min; + unsigned long sched_t_max; +} nir_driver_config_t; + +typedef struct { + float *accum_w; + float *accum_b; + size_t w_count; + size_t b_count; + nir_sgd_t sgd_w; + nir_sgd_t sgd_b; + nir_adam_t adam_w; + nir_adam_t adam_b; + int has_w; + int has_b; + int sgd_w_live; + int sgd_b_live; + int adam_w_live; + int adam_b_live; +} nir_driver_node_t; + +typedef struct { + nir_graph_t *graph; + nir_tape_t tape; + nir_driver_node_t *nodes; + size_t node_count; + nir_driver_config_t config; + nir_rng_t shuffle_rng; + nir_rng_t encode_rng; + unsigned long step; + int live; +} nir_driver_t; + +typedef struct { + float loss; + float accuracy; +} nir_driver_stats_t; + +/* Bind `graph` for training: validate `config`, init the tape, + * allocate one accumulator plus one SGD/Adam instance per + * weight/bias buffer, and seed the shuffle/encode streams. + * `graph` must outlive `driver`. Returns 0 with *error set. */ +int nir_driver_init( + nir_driver_t *driver, nir_graph_t *graph, + const nir_driver_config_t *config, char **error +); + +/* Free tape, accumulators, and optimizer state. Safe on a + * zeroed or already-freed driver. */ +void nir_driver_free(nir_driver_t *driver); + +/* One shuffled epoch over IDX buffers: encode each sample, + * tape forward, cross-entropy, tape backward, accumulate, + * mean-scale per batch, schedule the LR, step. Writes the + * epoch mean loss and argmax accuracy. Fails on shape, + * label-range, or non-finite errors with *error set. */ +int nir_driver_train_epoch( + nir_driver_t *driver, const uint8_t *pixels, + const uint8_t *labels, size_t count, size_t rows, + size_t cols, nir_driver_stats_t *stats, char **error +); + +/* Forward-only pass over the same buffers: no shuffle, no + * backward, no weight updates. Rate encoding still consumes + * the encode stream. */ +int nir_driver_eval( + nir_driver_t *driver, const uint8_t *pixels, + const uint8_t *labels, size_t count, size_t rows, + size_t cols, nir_driver_stats_t *stats, char **error +); + +/* Save weights plus optimizer state and the step counter. + * `topology` names the run (NULL omits it). */ +int nir_driver_save( + nir_driver_t *driver, const char *path, + const char *topology, char **error +); + +/* Load weights plus optimizer state and the step counter. + * Optimizer hyperparameters must already match the saver. */ +int nir_driver_load( + nir_driver_t *driver, const char *path, char **error +); + +/* Completed optimizer steps (the schedule counter). */ +unsigned long nir_driver_step(const nir_driver_t *driver); + +#endif /* NIR_C_RUNTIME_DRIVER_H */ diff --git a/include/nir_c_runtime/train.h b/include/nir_c_runtime/train.h index b6f7135..3e35453 100644 --- a/include/nir_c_runtime/train.h +++ b/include/nir_c_runtime/train.h @@ -3,11 +3,12 @@ * 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). + * then a torch.optim.Adam step). Single-step callers compose the + * backward pass from the layer backward ops in ops.h; the + * graph-level BPTT tape (tape.h) and the end-to-end driver + * (driver.h) compose it for whole graphs. 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 diff --git a/src/driver.c b/src/driver.c new file mode 100644 index 0000000..7425bf9 --- /dev/null +++ b/src/driver.c @@ -0,0 +1,324 @@ +#include "nir_c_runtime/driver.h" + +#include +#include +#include + +#include "nir_c_runtime/checkpoint.h" +#include "nir_c_runtime/data.h" +#include "nir_c_runtime/encode.h" +#include "nir_c_runtime/minibatch.h" +#include "nir_c_runtime/schedule.h" + +static void set_error(char **error, const char *msg) { + if (error == NULL || *error != NULL) { + return; + } + size_t n = strlen(msg) + 1; + *error = (char *)malloc(n); + if (*error != NULL) { + memcpy(*error, msg, n); + } +} + +static int node_weights( + nir_graph_t *g, size_t idx, float **w, size_t *nw, + float **b, size_t *nb +) { + nir_node_t *node = &g->nodes[idx]; + *w = NULL; + *nw = 0; + *b = NULL; + *nb = 0; + switch (node->kind) { + case NIR_NODE_AFFINE: + *w = node->params.affine.weight; + *nw = node->params.affine.out_features * + node->params.affine.in_features; + *b = node->params.affine.bias; + if (*b != NULL) { + *nb = node->params.affine.out_features; + } + return 1; + case NIR_NODE_LINEAR: + *w = node->params.linear.weight; + *nw = node->params.linear.out_features * + node->params.linear.in_features; + return 1; + case NIR_NODE_CONV2D: + *w = node->params.conv2d.weight; + *nw = node->params.conv2d.out_channels * + node->params.conv2d.in_channels * + node->params.conv2d.kernel_h * + node->params.conv2d.kernel_w; + *b = node->params.conv2d.bias; + if (*b != NULL) { + *nb = node->params.conv2d.out_channels; + } + return 1; + case NIR_NODE_CONV1D: + *w = node->params.conv1d.weight; + *nw = node->params.conv1d.out_channels * + node->params.conv1d.in_channels * + node->params.conv1d.kernel; + *b = node->params.conv1d.bias; + if (*b != NULL) { + *nb = node->params.conv1d.out_channels; + } + return 1; + default: + return 0; + } +} + +static int opt_hyper_ok( + const nir_driver_config_t *c, char **error +) { + if (!isfinite(c->lr) || c->lr < 0.0f || + !isfinite(c->weight_decay) || c->weight_decay < 0.0f) { + set_error(error, "driver needs lr and decay >= 0"); + return 0; + } + if (c->opt_kind == NIR_DRIVER_OPT_SGD) { + if (!isfinite(c->momentum) || c->momentum < 0.0f) { + set_error(error, "driver SGD needs momentum >= 0"); + return 0; + } + return 1; + } + if (c->opt_kind != NIR_DRIVER_OPT_ADAM) { + set_error(error, "driver needs SGD or Adam"); + return 0; + } + if (!isfinite(c->beta1) || c->beta1 < 0.0f || + c->beta1 >= 1.0f || !isfinite(c->beta2) || + c->beta2 < 0.0f || c->beta2 >= 1.0f || + !isfinite(c->eps) || c->eps <= 0.0f) { + set_error(error, "driver Adam needs betas, eps"); + return 0; + } + return 1; +} + +static int sched_ok( + const nir_driver_config_t *c, char **error +) { + switch (c->sched) { + case NIR_DRIVER_SCHED_NONE: + return 1; + case NIR_DRIVER_SCHED_STEP: + if (!isfinite(c->sched_gamma) || + c->sched_gamma < 0.0f || + c->sched_step_size == 0) { + set_error(error, "driver step schedule invalid"); + return 0; + } + return 1; + case NIR_DRIVER_SCHED_COSINE: + if (!isfinite(c->sched_eta_min) || + c->sched_eta_min < 0.0f || + c->sched_t_max == 0) { + set_error(error, "driver cosine schedule invalid"); + return 0; + } + return 1; + case NIR_DRIVER_SCHED_EXP: + if (!isfinite(c->sched_gamma) || + c->sched_gamma < 0.0f) { + set_error(error, "driver exp schedule invalid"); + return 0; + } + return 1; + default: + set_error(error, "driver schedule unknown"); + return 0; + } +} + +static int config_ok( + const nir_driver_config_t *c, char **error +) { + if (c == NULL || c->num_steps == 0 || + c->batch_size == 0) { + set_error(error, "driver needs steps and batch"); + return 0; + } + if (c->encode != NIR_DRIVER_ENCODE_REPEAT && + c->encode != NIR_DRIVER_ENCODE_RATE) { + set_error(error, "driver encode must be repeat/rate"); + return 0; + } + if (c->encode == NIR_DRIVER_ENCODE_RATE && + !isfinite(c->encode_gain)) { + set_error(error, "driver rate encode needs a gain"); + return 0; + } + float dummy = 0.0f; + if (!nir_surrogate_grad(c->surrogate, c->slope, 0.0f, + &dummy)) { + set_error(error, "driver surrogate/slope invalid"); + return 0; + } + return opt_hyper_ok(c, error) && sched_ok(c, error); +} + +static void free_nodes( + nir_driver_node_t *nodes, size_t count +) { + if (nodes == NULL) { + return; + } + for (size_t i = 0; i < count; i++) { + free(nodes[i].accum_w); + free(nodes[i].accum_b); + if (nodes[i].sgd_w_live) { + nir_sgd_free(&nodes[i].sgd_w); + } + if (nodes[i].sgd_b_live) { + nir_sgd_free(&nodes[i].sgd_b); + } + if (nodes[i].adam_w_live) { + nir_adam_free(&nodes[i].adam_w); + } + if (nodes[i].adam_b_live) { + nir_adam_free(&nodes[i].adam_b); + } + } + free(nodes); +} + +static int init_one_sgd( + nir_driver_t *d, size_t idx, char **error +) { + nir_driver_node_t *n = &d->nodes[idx]; + const nir_driver_config_t *c = &d->config; + if (!nir_sgd_init(&n->sgd_w, n->w_count, c->lr, + c->momentum, c->weight_decay, error)) { + return 0; + } + n->sgd_w_live = 1; + if (!n->has_b) { + return 1; + } + if (!nir_sgd_init(&n->sgd_b, n->b_count, c->lr, + c->momentum, c->weight_decay, error)) { + return 0; + } + n->sgd_b_live = 1; + return 1; +} + +static int init_one_adam( + nir_driver_t *d, size_t idx, char **error +) { + nir_driver_node_t *n = &d->nodes[idx]; + const nir_driver_config_t *c = &d->config; + if (!nir_adam_init(&n->adam_w, n->w_count, c->lr, + c->beta1, c->beta2, c->eps, c->weight_decay, + error)) { + return 0; + } + n->adam_w_live = 1; + if (!n->has_b) { + return 1; + } + if (!nir_adam_init(&n->adam_b, n->b_count, c->lr, + c->beta1, c->beta2, c->eps, c->weight_decay, + error)) { + return 0; + } + n->adam_b_live = 1; + return 1; +} + +static int init_nodes(nir_driver_t *d, char **error) { + size_t n = d->graph->node_count; + d->nodes = (nir_driver_node_t *)calloc( + n == 0 ? 1 : n, sizeof(nir_driver_node_t)); + if (d->nodes == NULL) { + set_error(error, "out of memory"); + return 0; + } + d->node_count = n; + for (size_t i = 0; i < n; i++) { + float *w = NULL; + float *b = NULL; + size_t nw = 0; + size_t nb = 0; + if (!node_weights(d->graph, i, &w, &nw, &b, + &nb)) { + continue; + } + nir_driver_node_t *node = &d->nodes[i]; + node->has_w = 1; + node->w_count = nw; + node->accum_w = (float *)calloc( + nw == 0 ? 1 : nw, sizeof(float)); + if (node->accum_w == NULL) { + set_error(error, "out of memory"); + return 0; + } + if (b != NULL && nb > 0) { + node->has_b = 1; + node->b_count = nb; + node->accum_b = (float *)calloc(nb, + sizeof(float)); + if (node->accum_b == NULL) { + set_error(error, "out of memory"); + return 0; + } + } + int ok = d->config.opt_kind == NIR_DRIVER_OPT_SGD + ? init_one_sgd(d, i, error) + : init_one_adam(d, i, error); + if (!ok) { + return 0; + } + } + return 1; +} + +int nir_driver_init( + nir_driver_t *driver, nir_graph_t *graph, + const nir_driver_config_t *config, char **error +) { + if (driver == NULL || graph == NULL) { + set_error(error, "driver needs a graph"); + return 0; + } + memset(driver, 0, sizeof(*driver)); + if (!config_ok(config, error)) { + return 0; + } + driver->graph = graph; + driver->config = *config; + driver->step = 0; + if (!nir_tape_init(&driver->tape, graph, + config->num_steps, config->surrogate, + config->slope, error)) { + return 0; + } + if (!init_nodes(driver, error)) { + nir_tape_free(&driver->tape); + free_nodes(driver->nodes, driver->node_count); + memset(driver, 0, sizeof(*driver)); + return 0; + } + nir_rng_seed(&driver->shuffle_rng, config->seed); + nir_rng_seed(&driver->encode_rng, config->seed + 1); + driver->live = 1; + return 1; +} + +void nir_driver_free(nir_driver_t *driver) { + if (driver == NULL) { + return; + } + nir_tape_free(&driver->tape); + free_nodes(driver->nodes, driver->node_count); + memset(driver, 0, sizeof(*driver)); +} + +unsigned long nir_driver_step(const nir_driver_t *driver) { + return driver == NULL ? 0 : driver->step; +} diff --git a/src/driver_loop.c b/src/driver_loop.c new file mode 100644 index 0000000..4a549d6 --- /dev/null +++ b/src/driver_loop.c @@ -0,0 +1,646 @@ +#include "nir_c_runtime/driver.h" + +#include +#include +#include + +#include "nir_c_runtime/checkpoint.h" +#include "nir_c_runtime/data.h" +#include "nir_c_runtime/encode.h" +#include "nir_c_runtime/minibatch.h" +#include "nir_c_runtime/schedule.h" + +static void set_error(char **error, const char *msg) { + if (error == NULL || *error != NULL) { + return; + } + size_t n = strlen(msg) + 1; + *error = (char *)malloc(n); + if (*error != NULL) { + memcpy(*error, msg, n); + } +} + +static void free_frames( + nir_tensor_t *frames, size_t count +) { + if (frames == NULL) { + return; + } + for (size_t i = 0; i < count; i++) { + nir_tensor_free(&frames[i]); + } +} + +static int frames_from_floats( + const float *flat, size_t n, size_t rows, size_t cols, + nir_tensor_t *frames, size_t steps, char **error +) { + size_t shape[3] = {1, rows, cols}; + for (size_t t = 0; t < steps; t++) { + frames[t] = nir_tensor_alloc(shape, 3); + if (frames[t].data == NULL) { + free_frames(frames, t); + set_error(error, "out of memory"); + return 0; + } + memcpy(frames[t].data, flat + t * n, + n * sizeof(float)); + } + return 1; +} + +static int frames_from_spikes( + const uint8_t *spikes, size_t n, size_t rows, + size_t cols, nir_tensor_t *frames, size_t steps, + char **error +) { + size_t shape[3] = {1, rows, cols}; + for (size_t t = 0; t < steps; t++) { + frames[t] = nir_tensor_alloc(shape, 3); + if (frames[t].data == NULL) { + free_frames(frames, t); + set_error(error, "out of memory"); + return 0; + } + for (size_t i = 0; i < n; i++) { + frames[t].data[i] = + spikes[t * n + i] ? 1.0f : 0.0f; + } + } + return 1; +} + +static int encode_repeat_frames( + nir_driver_t *d, const float *norm, size_t n, + size_t rows, size_t cols, nir_tensor_t *frames, + char **error +) { + size_t steps = d->config.num_steps; + float *flat = (float *)malloc( + steps * n * sizeof(float)); + if (flat == NULL) { + set_error(error, "out of memory"); + return 0; + } + int ok = nir_encode_repeat(norm, n, steps, flat, + error); + if (ok) { + ok = frames_from_floats(flat, n, rows, cols, + frames, steps, error); + } + free(flat); + return ok; +} + +static int encode_rate_frames( + nir_driver_t *d, const float *norm, size_t n, + size_t rows, size_t cols, nir_tensor_t *frames, + char **error +) { + size_t steps = d->config.num_steps; + uint8_t *spikes = (uint8_t *)malloc(steps * n); + if (spikes == NULL) { + set_error(error, "out of memory"); + return 0; + } + int ok = nir_encode_rate(norm, n, steps, + d->config.encode_gain, &d->encode_rng, spikes, + error); + if (ok) { + ok = frames_from_spikes(spikes, n, rows, cols, + frames, steps, error); + } + free(spikes); + return ok; +} + +static int encode_sample( + nir_driver_t *d, const uint8_t *pixels, size_t rows, + size_t cols, size_t index, nir_tensor_t *frames, + char **error +) { + size_t n = rows * cols; + const uint8_t *raw = pixels + index * n; + float *norm = (float *)malloc(n * sizeof(float)); + if (norm == NULL) { + set_error(error, "out of memory"); + return 0; + } + nir_data_normalize(raw, norm, n); + int ok = 0; + if (d->config.encode == NIR_DRIVER_ENCODE_REPEAT) { + ok = encode_repeat_frames(d, norm, n, rows, + cols, frames, error); + } else { + ok = encode_rate_frames(d, norm, n, rows, cols, + frames, error); + } + free(norm); + return ok; +} + +static size_t argmax(const nir_tensor_t *t) { + size_t best = 0; + for (size_t i = 1; i < t->size; i++) { + if (t->data[i] > t->data[best]) { + best = i; + } + } + return best; +} + +static int accumulate_tape_grads( + nir_driver_t *d, char **error +) { + for (size_t i = 0; i < d->node_count; i++) { + nir_driver_node_t *node = &d->nodes[i]; + if (!node->has_w) { + continue; + } + const nir_tensor_t *gw = nir_tape_weight_grad( + &d->tape, i); + if (gw == NULL || gw->size != node->w_count) { + set_error(error, "driver missing weight grad"); + return 0; + } + if (!nir_minibatch_accumulate(node->accum_w, + gw->data, node->w_count)) { + set_error(error, "driver grad not finite"); + return 0; + } + if (!node->has_b) { + continue; + } + const nir_tensor_t *gb = nir_tape_bias_grad( + &d->tape, i); + if (gb == NULL || gb->size != node->b_count) { + set_error(error, "driver missing bias grad"); + return 0; + } + if (!nir_minibatch_accumulate(node->accum_b, + gb->data, node->b_count)) { + set_error(error, "driver grad not finite"); + return 0; + } + } + return 1; +} + +static int forward_loss( + nir_driver_t *d, nir_tensor_t *frames, size_t label, + nir_tensor_t *readout, nir_tensor_t *grad, + float *loss, char **error +) { + memset(readout, 0, sizeof(*readout)); + if (!nir_tape_forward(&d->tape, frames, readout, + error)) { + return 0; + } + size_t shape[1] = {readout->size}; + *grad = nir_tensor_alloc(shape, 1); + if (grad->data == NULL) { + nir_tensor_free(readout); + set_error(error, "out of memory"); + return 0; + } + if (!nir_cross_entropy(readout, label, loss, grad)) { + nir_tensor_free(readout); + nir_tensor_free(grad); + set_error(error, "driver loss failed"); + return 0; + } + return 1; +} + +static int train_one( + nir_driver_t *d, const uint8_t *pixels, + size_t sample, size_t rows, size_t cols, + size_t label, float *loss, int *correct, + char **error +) { + size_t steps = d->config.num_steps; + nir_tensor_t *frames = (nir_tensor_t *)calloc( + steps, sizeof(nir_tensor_t)); + if (frames == NULL) { + set_error(error, "out of memory"); + return 0; + } + int ok = 0; + nir_tensor_t readout; + nir_tensor_t grad; + memset(&grad, 0, sizeof(grad)); + if (!encode_sample(d, pixels, rows, cols, sample, + frames, error)) { + goto done; + } + if (!forward_loss(d, frames, label, &readout, &grad, + loss, error)) { + goto done; + } + *correct = argmax(&readout) == label; + if (!nir_tape_backward(&d->tape, &grad, error)) { + nir_tensor_free(&readout); + nir_tensor_free(&grad); + goto done; + } + ok = accumulate_tape_grads(d, error); + nir_tensor_free(&readout); + nir_tensor_free(&grad); +done: + free_frames(frames, steps); + free(frames); + return ok; +} + +static void zero_accums(nir_driver_t *d) { + for (size_t i = 0; i < d->node_count; i++) { + nir_driver_node_t *n = &d->nodes[i]; + if (n->has_w) { + nir_minibatch_zero(n->accum_w, n->w_count); + } + if (n->has_b) { + nir_minibatch_zero(n->accum_b, n->b_count); + } + } +} + +static int scale_accums( + nir_driver_t *d, size_t batch, char **error +) { + float scale = 1.0f / (float)batch; + for (size_t i = 0; i < d->node_count; i++) { + nir_driver_node_t *n = &d->nodes[i]; + if (n->has_w && + !nir_minibatch_scale(n->accum_w, n->w_count, + scale)) { + set_error(error, "driver scale failed"); + return 0; + } + if (n->has_b && + !nir_minibatch_scale(n->accum_b, n->b_count, + scale)) { + set_error(error, "driver scale failed"); + return 0; + } + } + return 1; +} + +static int step_node( + nir_driver_t *d, size_t idx, char **error +) { + nir_driver_node_t *n = &d->nodes[idx]; + if (!n->has_w) { + return 1; + } + const char *name = d->graph->nodes[idx].name; + const float *cw = NULL; + const float *cb = NULL; + size_t nw = 0; + size_t nb = 0; + if (!nir_graph_get_weights(d->graph, name, &cw, &nw, + &cb, &nb)) { + set_error(error, "driver weights missing"); + return 0; + } + float *w = (float *)cw; + float *b = (float *)cb; + int ok = 1; + if (d->config.opt_kind == NIR_DRIVER_OPT_SGD) { + ok = nir_sgd_step(&n->sgd_w, w, n->accum_w, nw); + if (ok && n->has_b) { + ok = nir_sgd_step(&n->sgd_b, b, n->accum_b, + nb); + } + } else { + ok = nir_adam_step(&n->adam_w, w, n->accum_w, + nw); + if (ok && n->has_b) { + ok = nir_adam_step(&n->adam_b, b, n->accum_b, + nb); + } + } + if (!ok) { + set_error(error, "driver optimizer step failed"); + } + return ok; +} + +static int current_lr( + nir_driver_t *d, float *lr, char **error +) { + const nir_driver_config_t *c = &d->config; + int ok = 1; + switch (c->sched) { + case NIR_DRIVER_SCHED_NONE: + *lr = c->lr; + break; + case NIR_DRIVER_SCHED_STEP: + ok = nir_lr_step_decay(c->lr, c->sched_gamma, + c->sched_step_size, d->step, lr); + break; + case NIR_DRIVER_SCHED_COSINE: + ok = nir_lr_cosine(c->lr, c->sched_eta_min, + c->sched_t_max, d->step, lr); + break; + default: + ok = nir_lr_exponential(c->lr, + c->sched_gamma, d->step, lr); + break; + } + if (!ok) { + set_error(error, "driver schedule failed"); + } + return ok; +} + +static void write_lr(nir_driver_t *d, float lr) { + for (size_t i = 0; i < d->node_count; i++) { + nir_driver_node_t *n = &d->nodes[i]; + if (n->sgd_w_live) { + n->sgd_w.lr = lr; + } + if (n->sgd_b_live) { + n->sgd_b.lr = lr; + } + if (n->adam_w_live) { + n->adam_w.lr = lr; + } + if (n->adam_b_live) { + n->adam_b.lr = lr; + } + } +} + +static int schedule_and_step( + nir_driver_t *d, char **error +) { + float lr = 0.0f; + if (!current_lr(d, &lr, error)) { + return 0; + } + write_lr(d, lr); + for (size_t i = 0; i < d->node_count; i++) { + if (!step_node(d, i, error)) { + return 0; + } + } + d->step++; + return 1; +} + +static int epoch_args_ok( + nir_driver_t *d, const uint8_t *px, + const uint8_t *lb, size_t count, size_t rows, + size_t cols, nir_driver_stats_t *st, char **error +) { + if (d == NULL || !d->live || px == NULL || + lb == NULL || count == 0 || rows == 0 || + cols == 0 || st == NULL) { + set_error(error, "driver epoch needs data"); + return 0; + } + return 1; +} + +static int run_batch( + nir_driver_t *d, const uint8_t *px, + const uint8_t *lb, const size_t *order, + size_t begin, size_t end, size_t rows, size_t cols, + double *loss_sum, size_t *correct, char **error +) { + zero_accums(d); + for (size_t k = begin; k < end; k++) { + size_t sample = order[k]; + float loss = 0.0f; + int hit = 0; + if (!train_one(d, px, sample, rows, cols, + (size_t)lb[sample], &loss, &hit, + error)) { + return 0; + } + *loss_sum += (double)loss; + if (hit) { + (*correct)++; + } + } + if (!scale_accums(d, end - begin, error)) { + return 0; + } + return schedule_and_step(d, error); +} + +int nir_driver_train_epoch( + nir_driver_t *driver, const uint8_t *pixels, + const uint8_t *labels, size_t count, size_t rows, + size_t cols, nir_driver_stats_t *stats, char **error +) { + if (!epoch_args_ok(driver, pixels, labels, count, + rows, cols, stats, error)) { + return 0; + } + size_t *order = (size_t *)malloc( + count * sizeof(size_t)); + if (order == NULL) { + set_error(error, "out of memory"); + return 0; + } + nir_sampler_arange(order, count); + if (!nir_sampler_shuffle(&driver->shuffle_rng, order, + count)) { + free(order); + set_error(error, "driver shuffle failed"); + return 0; + } + double loss_sum = 0.0; + size_t correct = 0; + size_t batch = driver->config.batch_size; + int ok = 1; + for (size_t s = 0; s < count; s += batch) { + size_t e = s + batch; + if (e > count) { + e = count; + } + if (!run_batch(driver, pixels, labels, order, s, + e, rows, cols, &loss_sum, &correct, + error)) { + ok = 0; + break; + } + } + free(order); + if (!ok) { + return 0; + } + stats->loss = (float)(loss_sum / (double)count); + stats->accuracy = (float)correct / (float)count; + return 1; +} + +static int eval_one( + nir_driver_t *d, const uint8_t *pixels, + size_t sample, size_t rows, size_t cols, + size_t label, float *loss, int *correct, + char **error +) { + size_t steps = d->config.num_steps; + nir_tensor_t *frames = (nir_tensor_t *)calloc( + steps, sizeof(nir_tensor_t)); + if (frames == NULL) { + set_error(error, "out of memory"); + return 0; + } + int ok = 0; + nir_tensor_t readout; + nir_tensor_t grad; + memset(&grad, 0, sizeof(grad)); + if (!encode_sample(d, pixels, rows, cols, sample, + frames, error)) { + goto done; + } + if (!forward_loss(d, frames, label, &readout, &grad, + loss, error)) { + goto done; + } + *correct = argmax(&readout) == label; + ok = 1; + nir_tensor_free(&readout); + nir_tensor_free(&grad); +done: + free_frames(frames, steps); + free(frames); + return ok; +} + +int nir_driver_eval( + nir_driver_t *driver, const uint8_t *pixels, + const uint8_t *labels, size_t count, size_t rows, + size_t cols, nir_driver_stats_t *stats, char **error +) { + if (!epoch_args_ok(driver, pixels, labels, count, + rows, cols, stats, error)) { + return 0; + } + double loss_sum = 0.0; + size_t correct = 0; + for (size_t i = 0; i < count; i++) { + float loss = 0.0f; + int hit = 0; + if (!eval_one(driver, pixels, i, rows, cols, + (size_t)labels[i], &loss, &hit, + error)) { + return 0; + } + loss_sum += (double)loss; + if (hit) { + correct++; + } + } + stats->loss = (float)(loss_sum / (double)count); + stats->accuracy = (float)correct / (float)count; + return 1; +} + +static size_t count_trained(const nir_driver_t *d) { + size_t n = 0; + for (size_t i = 0; i < d->node_count; i++) { + if (d->nodes[i].has_w) { + n++; + } + } + return n; +} + +static void fill_opt( + nir_driver_t *d, size_t idx, nir_checkpoint_opt_t *o +) { + nir_driver_node_t *n = &d->nodes[idx]; + memset(o, 0, sizeof(*o)); + o->node = d->graph->nodes[idx].name; + if (d->config.opt_kind == NIR_DRIVER_OPT_SGD) { + o->kind = NIR_CHECKPOINT_OPT_SGD; + o->sgd_weight = &n->sgd_w; + o->sgd_bias = n->has_b ? &n->sgd_b : NULL; + } else { + o->kind = NIR_CHECKPOINT_OPT_ADAM; + o->adam_weight = &n->adam_w; + o->adam_bias = n->has_b ? &n->adam_b : NULL; + } +} + +static nir_checkpoint_opt_t *build_opts( + nir_driver_t *d, size_t *count, char **error +) { + size_t n = count_trained(d); + *count = n; + if (n == 0) { + return NULL; + } + nir_checkpoint_opt_t *opts = + (nir_checkpoint_opt_t *)calloc(n, + sizeof(nir_checkpoint_opt_t)); + if (opts == NULL) { + set_error(error, "out of memory"); + return NULL; + } + size_t k = 0; + for (size_t i = 0; i < d->node_count; i++) { + if (d->nodes[i].has_w) { + fill_opt(d, i, &opts[k++]); + } + } + return opts; +} + +int nir_driver_save( + nir_driver_t *driver, const char *path, + const char *topology, char **error +) { + if (driver == NULL || !driver->live || + path == NULL) { + set_error(error, "driver save needs a path"); + return 0; + } + size_t count = 0; + nir_checkpoint_opt_t *opts = build_opts(driver, + &count, error); + if (count > 0 && opts == NULL) { + return 0; + } + nir_checkpoint_meta_t meta; + memset(&meta, 0, sizeof(meta)); + meta.topology = (char *)topology; + meta.seed = driver->config.seed; + meta.step = driver->step; + int ok = nir_checkpoint_save_with_opt(path, + driver->graph, &meta, opts, count, error); + free(opts); + return ok; +} + +int nir_driver_load( + nir_driver_t *driver, const char *path, char **error +) { + if (driver == NULL || !driver->live || + path == NULL) { + set_error(error, "driver load needs a path"); + return 0; + } + size_t count = 0; + nir_checkpoint_opt_t *opts = build_opts(driver, + &count, error); + if (count > 0 && opts == NULL) { + return 0; + } + nir_checkpoint_meta_t meta; + memset(&meta, 0, sizeof(meta)); + int ok = nir_checkpoint_load_with_opt(path, + driver->graph, &meta, opts, count, error); + if (ok) { + driver->step = meta.step; + } + nir_checkpoint_meta_free(&meta); + free(opts); + return ok; +} diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 770627a..8a40464 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -28,6 +28,7 @@ nir_add_test(test_data) nir_add_test(test_checkpoint) nir_add_test(test_minibatch) nir_add_test(test_tape) +nir_add_test(test_driver) # 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_driver.c b/tests/test_driver.c new file mode 100644 index 0000000..9deb31b --- /dev/null +++ b/tests/test_driver.c @@ -0,0 +1,303 @@ +/* End-to-end driver: loss falls on the tiny fixture, + * checkpoints resume bit-identically, and the SGD/rate/ + * schedule paths run cleanly. */ +#include +#include +#include +#include +#include + +#include "nir_c_runtime/checkpoint.h" +#include "nir_c_runtime/data.h" +#include "nir_c_runtime/driver.h" +#include "nir_c_runtime/graph.h" +#include "nir_c_runtime/json.h" +#include "test_util.h" + +static nir_graph_t load_graph(const char *path) { + nir_graph_t graph; + memset(&graph, 0, sizeof(graph)); + char *text = read_file(path); + char *error = NULL; + json_value_t *json = json_parse(text, &error); + CHECK(error == NULL); + CHECK(nir_graph_load(json, &graph, &error)); + CHECK(error == NULL); + json_free(json); + free(text); + return graph; +} + +static void base_config(nir_driver_config_t *c) { + memset(c, 0, sizeof(*c)); + c->num_steps = 4; + c->surrogate = NIR_SURROGATE_FAST_SIGMOID; + c->slope = 25.0f; + c->encode = NIR_DRIVER_ENCODE_REPEAT; + c->encode_gain = 1.0f; + c->batch_size = 2; + c->seed = 7; + c->opt_kind = NIR_DRIVER_OPT_ADAM; + c->lr = 0.01f; + c->weight_decay = 0.0f; + c->beta1 = 0.9f; + c->beta2 = 0.999f; + c->eps = 1e-8f; + c->sched = NIR_DRIVER_SCHED_NONE; +} + +static void load_tiny( + uint8_t **px, uint8_t **lb, size_t *count, + size_t *rows, size_t *cols +) { + char *error = NULL; + size_t nlabels = 0; + CHECK(nir_idx_read_images( + "fixtures/mnist_tiny_images.idx", px, count, + rows, cols, &error)); + CHECK(error == NULL); + CHECK(nir_idx_read_labels( + "fixtures/mnist_tiny_labels.idx", lb, &nlabels, + &error)); + CHECK(error == NULL); + CHECK(nlabels == *count); +} + +static void temp_path(const char *name, char *out) { + const char *tmp = getenv("TMPDIR"); + if (tmp == NULL) { + tmp = "/tmp"; + } + sprintf(out, "%s/nir-driver-test-%s", tmp, name); +} + +static void test_train_decreases(void) { + nir_graph_t graph = load_graph( + "../examples/mnist_tiny_graph.json"); + uint8_t *px = NULL; + uint8_t *lb = NULL; + size_t count = 0; + size_t rows = 0; + size_t cols = 0; + load_tiny(&px, &lb, &count, &rows, &cols); + nir_driver_config_t config; + base_config(&config); + nir_driver_t driver; + memset(&driver, 0, sizeof(driver)); + char *error = NULL; + CHECK(nir_driver_init(&driver, &graph, &config, + &error)); + CHECK(error == NULL); + nir_driver_stats_t before; + CHECK(nir_driver_eval(&driver, px, lb, count, rows, + cols, &before, &error)); + CHECK(error == NULL); + CHECK(isfinite(before.loss)); + for (int e = 0; e < 20; e++) { + nir_driver_stats_t st; + CHECK(nir_driver_train_epoch(&driver, px, lb, + count, rows, cols, &st, &error)); + CHECK(error == NULL); + CHECK(isfinite(st.loss)); + } + nir_driver_stats_t after; + CHECK(nir_driver_eval(&driver, px, lb, count, rows, + cols, &after, &error)); + CHECK(error == NULL); + CHECK(isfinite(after.loss)); + /* Modest bound: 20 Adam epochs beat init by 0.5. */ + CHECK(after.loss + 0.5f < before.loss); + CHECK(nir_driver_step(&driver) == 40); + nir_driver_free(&driver); + free(px); + free(lb); + nir_graph_free(&graph); +} + +static void test_checkpoint_resume(void) { + nir_graph_t graph = load_graph( + "../examples/mnist_tiny_graph.json"); + uint8_t *px = NULL; + uint8_t *lb = NULL; + size_t count = 0; + size_t rows = 0; + size_t cols = 0; + load_tiny(&px, &lb, &count, &rows, &cols); + nir_driver_config_t config; + base_config(&config); + nir_driver_t driver; + memset(&driver, 0, sizeof(driver)); + char *error = NULL; + CHECK(nir_driver_init(&driver, &graph, &config, + &error)); + for (int e = 0; e < 5; e++) { + nir_driver_stats_t st; + CHECK(nir_driver_train_epoch(&driver, px, lb, + count, rows, cols, &st, &error)); + } + nir_driver_stats_t trained; + CHECK(nir_driver_eval(&driver, px, lb, count, rows, + cols, &trained, &error)); + char path[512]; + temp_path("resume.json", path); + CHECK(nir_driver_save(&driver, path, "mnist_tiny", + &error)); + CHECK(error == NULL); + unsigned long want_step = nir_driver_step(&driver); + CHECK(want_step == 10); + nir_driver_free(&driver); + nir_graph_free(&graph); + nir_graph_t graph2 = load_graph( + "../examples/mnist_tiny_graph.json"); + nir_driver_t driver2; + memset(&driver2, 0, sizeof(driver2)); + CHECK(nir_driver_init(&driver2, &graph2, &config, + &error)); + CHECK(nir_driver_load(&driver2, path, &error)); + CHECK(error == NULL); + CHECK(nir_driver_step(&driver2) == want_step); + nir_driver_stats_t resumed; + CHECK(nir_driver_eval(&driver2, px, lb, count, rows, + cols, &resumed, &error)); + CHECK(resumed.loss == trained.loss); + CHECK(resumed.accuracy == trained.accuracy); + nir_driver_free(&driver2); + nir_graph_free(&graph2); + free(px); + free(lb); +} + +static void test_sgd_rate_sched(void) { + nir_graph_t graph = load_graph( + "../examples/mnist_tiny_graph.json"); + uint8_t *px = NULL; + uint8_t *lb = NULL; + size_t count = 0; + size_t rows = 0; + size_t cols = 0; + load_tiny(&px, &lb, &count, &rows, &cols); + char *error = NULL; + nir_driver_config_t sgd; + base_config(&sgd); + sgd.opt_kind = NIR_DRIVER_OPT_SGD; + sgd.momentum = 0.9f; + nir_driver_t d1; + memset(&d1, 0, sizeof(d1)); + CHECK(nir_driver_init(&d1, &graph, &sgd, &error)); + for (int e = 0; e < 2; e++) { + nir_driver_stats_t st; + CHECK(nir_driver_train_epoch(&d1, px, lb, + count, rows, cols, &st, &error)); + CHECK(isfinite(st.loss)); + } + nir_driver_free(&d1); + nir_driver_config_t rate; + base_config(&rate); + rate.encode = NIR_DRIVER_ENCODE_RATE; + nir_driver_t d2; + memset(&d2, 0, sizeof(d2)); + CHECK(nir_driver_init(&d2, &graph, &rate, &error)); + for (int e = 0; e < 2; e++) { + nir_driver_stats_t st; + CHECK(nir_driver_train_epoch(&d2, px, lb, + count, rows, cols, &st, &error)); + CHECK(isfinite(st.loss)); + } + nir_driver_free(&d2); + nir_driver_config_t sched; + base_config(&sched); + sched.sched = NIR_DRIVER_SCHED_STEP; + sched.sched_gamma = 0.9f; + sched.sched_step_size = 5; + nir_driver_t d3; + memset(&d3, 0, sizeof(d3)); + CHECK(nir_driver_init(&d3, &graph, &sched, &error)); + nir_driver_stats_t st; + CHECK(nir_driver_train_epoch(&d3, px, lb, count, + rows, cols, &st, &error)); + CHECK(isfinite(st.loss)); + nir_driver_free(&d3); + nir_graph_free(&graph); + free(px); + free(lb); +} + +static void test_errors(void) { + nir_graph_t graph = load_graph( + "../examples/mnist_tiny_graph.json"); + uint8_t *px = NULL; + uint8_t *lb = NULL; + size_t count = 0; + size_t rows = 0; + size_t cols = 0; + load_tiny(&px, &lb, &count, &rows, &cols); + char *error = NULL; + nir_driver_config_t bad; + base_config(&bad); + nir_driver_t driver; + memset(&driver, 0, sizeof(driver)); + bad.num_steps = 0; + CHECK(!nir_driver_init(&driver, &graph, &bad, + &error)); + CHECK(error != NULL); + free(error); + error = NULL; + base_config(&bad); + bad.batch_size = 0; + CHECK(!nir_driver_init(&driver, &graph, &bad, + &error)); + free(error); + error = NULL; + base_config(&bad); + bad.lr = -1.0f; + CHECK(!nir_driver_init(&driver, &graph, &bad, + &error)); + free(error); + error = NULL; + base_config(&bad); + bad.slope = 0.0f; + CHECK(!nir_driver_init(&driver, &graph, &bad, + &error)); + free(error); + error = NULL; + nir_driver_config_t good; + base_config(&good); + CHECK(nir_driver_init(&driver, &graph, &good, + &error)); + CHECK(error == NULL); + nir_driver_stats_t st; + CHECK(!nir_driver_train_epoch(&driver, NULL, lb, + count, rows, cols, &st, &error)); + free(error); + error = NULL; + uint8_t bad_labels[3] = {0, 10, 1}; + CHECK(!nir_driver_train_epoch(&driver, px, + bad_labels, count, rows, cols, &st, &error)); + free(error); + error = NULL; + CHECK(!nir_driver_eval(&driver, px, bad_labels, + count, rows, cols, &st, &error)); + free(error); + error = NULL; + CHECK(!nir_driver_save(&driver, NULL, "t", &error)); + free(error); + error = NULL; + CHECK(!nir_driver_load(&driver, NULL, &error)); + free(error); + error = NULL; + CHECK(nir_driver_step(NULL) == 0); + nir_driver_free(&driver); + nir_driver_free(&driver); + nir_driver_free(NULL); + nir_graph_free(&graph); + free(px); + free(lb); +} + +int main(void) { + test_train_decreases(); + test_checkpoint_resume(); + test_sgd_rate_sched(); + test_errors(); + return 0; +}