diff --git a/CMakeLists.txt b/CMakeLists.txt index fd315b5..4ab1aa7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -26,6 +26,7 @@ add_library(nir_c_runtime src/encode.c src/data.c src/checkpoint.c + src/minibatch.c ) target_include_directories(nir_c_runtime PUBLIC include) target_link_libraries(nir_c_runtime PUBLIC m) diff --git a/include/nir_c_runtime/data.h b/include/nir_c_runtime/data.h index 185072c..83c5ebe 100644 --- a/include/nir_c_runtime/data.h +++ b/include/nir_c_runtime/data.h @@ -8,7 +8,9 @@ * 1/255 and the Normalize is an identity -- so C decodes each * pixel as v / 255.0f. Multi-byte samples land in `nir_tensor_t` * frames (batch-1, one tensor per sample); batching stays the - * caller's loop until the tensor ABI grows a batch axis. + * caller's loop with the mean-reduction convention in + * minibatch.h, which also records the #42 decision against a + * batch axis on the tensor ABI. */ #ifndef NIR_C_RUNTIME_DATA_H #define NIR_C_RUNTIME_DATA_H diff --git a/include/nir_c_runtime/loss.h b/include/nir_c_runtime/loss.h index b727896..db64256 100644 --- a/include/nir_c_runtime/loss.h +++ b/include/nir_c_runtime/loss.h @@ -25,8 +25,8 @@ * count gradient -- until the BPTT tape (#39) owns that. * * Both keep the single-sample-first convention: 1-D vectors, - * mean reduction for MSE, batch reduction left to the - * mini-batch issue (#42). + * mean reduction for MSE, batch reduction via the mean + * convention in minibatch.h. */ #ifndef NIR_C_RUNTIME_LOSS_H #define NIR_C_RUNTIME_LOSS_H diff --git a/include/nir_c_runtime/minibatch.h b/include/nir_c_runtime/minibatch.h new file mode 100644 index 0000000..d2679dd --- /dev/null +++ b/include/nir_c_runtime/minibatch.h @@ -0,0 +1,80 @@ +/* Mini-batch training over batch-1 tensors (#42). + * + * Reduction convention (this is the documented contract; train.h + * and data.h point here): the caller loops over one sample at a + * time and reduces with a mean, matching torch's default + * `reduction='mean'`: + * + * batch loss = mean(sample losses) + * batch grad = mean(sample grads) + * + * Per batch the caller zeroes one accumulator per parameter + * buffer (`nir_minibatch_zero`), adds each sample's gradient + * (`nir_minibatch_accumulate`), scales the sums by 1/B + * (`nir_minibatch_scale`), and only then takes the optimizer + * step. Losses reduce the same way through + * `nir_minibatch_mean_loss`. The optimizers in train.h always + * see the mean gradient -- scale exactly once per batch. + * + * Within one sample the single-sample losses keep their own + * rules: cross-entropy (train.h) and spike-count loss (loss.h) + * are per-logit-vector, MSE (loss.h) already means over its + * elements. Those per-sample values are what the batch mean + * averages; nothing here re-weights them. + * + * Batch-axis decision (#42, decided): no batch axis on + * `nir_tensor_t`. The accumulation loop above keeps the tensor + * ABI and every inference call site unchanged, and the extra + * adds are exact and cheap next to a forward/backward pass. + * Revisit only with measured evidence that the loop -- not the + * math under it -- is the bottleneck or the bug farm. + * + * Epoch shuffling is Fisher-Yates over an index array + * (`nir_sampler_arange` + `nir_sampler_shuffle`), driven by the + * seeded `nir_rng_t` stream so an epoch order replays exactly + * from its seed. + */ +#ifndef NIR_C_RUNTIME_MINIBATCH_H +#define NIR_C_RUNTIME_MINIBATCH_H + +#include + +#include "nir_c_runtime/rng.h" + +/* Zero an accumulator of `n` floats. `accum` must hold `n` + * floats; `n` may be 0. */ +void nir_minibatch_zero(float *accum, size_t n); + +/* Add `grad` into `accum` elementwise. Both buffers must hold + * `n` floats with `n` nonzero. Fails (return 0, `accum` + * untouched) on NULL buffers or non-finite data. */ +int nir_minibatch_accumulate( + float *accum, const float *grad, size_t n +); + +/* Scale `accum` by `scale` elementwise -- the 1/B step that + * turns accumulated sums into the mean gradient. Fails + * (return 0, `accum` untouched) on a NULL buffer, `n` of 0, + * or non-finite scale or data. */ +int nir_minibatch_scale(float *accum, size_t n, float scale); + +/* Mean of `n` per-sample losses, accumulated in float64 and + * rounded once. Fails (return 0) on NULL buffers, `n` of 0, + * or a non-finite entry. */ +int nir_minibatch_mean_loss( + const float *losses, size_t n, float *mean +); + +/* Fill `indices` with 0..n-1. `indices` must hold `n` + * entries; `n` may be 0. */ +void nir_sampler_arange(size_t *indices, size_t n); + +/* Fisher-Yates shuffle of `indices` in place, drawing one + * `nir_rng_next_u64` per position from `i = n-1` down to 1 + * with `j = draw % (i + 1)`. A seed therefore replays the + * same permutation on any platform. Fails (return 0) on a + * NULL `rng`, or a NULL `indices` with `n` nonzero; `n` of + * 0 or 1 is a successful no-op. */ +int nir_sampler_shuffle(nir_rng_t *rng, size_t *indices, size_t n); + +#endif /* NIR_C_RUNTIME_MINIBATCH_H */ diff --git a/include/nir_c_runtime/train.h b/include/nir_c_runtime/train.h index 95088fe..b6f7135 100644 --- a/include/nir_c_runtime/train.h +++ b/include/nir_c_runtime/train.h @@ -28,7 +28,8 @@ int nir_surrogate_from_name( /* 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. */ + * grad[target] -= 1. Batch reduction is the caller's loop with + * the mean convention in minibatch.h. */ int nir_cross_entropy( const nir_tensor_t *logits, size_t target, float *loss, nir_tensor_t *grad diff --git a/src/minibatch.c b/src/minibatch.c new file mode 100644 index 0000000..c0f8daa --- /dev/null +++ b/src/minibatch.c @@ -0,0 +1,75 @@ +#include "nir_c_runtime/minibatch.h" + +#include + +void nir_minibatch_zero(float *accum, size_t n) { + for (size_t i = 0; i < n; i++) { + accum[i] = 0.0f; + } +} + +static int check_finite(const float *data, size_t n) { + for (size_t i = 0; i < n; i++) { + if (!isfinite(data[i])) return 0; + } + return 1; +} + +int nir_minibatch_accumulate( + float *accum, const float *grad, size_t n +) { + if (accum == NULL || grad == NULL || n == 0 || + !check_finite(accum, n) || !check_finite(grad, n)) { + return 0; + } + for (size_t i = 0; i < n; i++) { + accum[i] += grad[i]; + } + return 1; +} + +int nir_minibatch_scale(float *accum, size_t n, float scale) { + if (accum == NULL || n == 0 || !isfinite(scale) || + !check_finite(accum, n)) { + return 0; + } + for (size_t i = 0; i < n; i++) { + accum[i] *= scale; + } + return 1; +} + +int nir_minibatch_mean_loss( + const float *losses, size_t n, float *mean +) { + if (losses == NULL || mean == NULL || n == 0 || + !check_finite(losses, n)) { + return 0; + } + double total = 0.0; + for (size_t i = 0; i < n; i++) { + total += (double)losses[i]; + } + *mean = (float)(total / (double)n); + return 1; +} + +void nir_sampler_arange(size_t *indices, size_t n) { + for (size_t i = 0; i < n; i++) { + indices[i] = i; + } +} + +int nir_sampler_shuffle(nir_rng_t *rng, size_t *indices, size_t n) { + if (rng == NULL || (indices == NULL && n > 0)) { + return 0; + } + for (size_t i = n; i > 1; i--) { + size_t j = (size_t)(nir_rng_next_u64(rng) % i); + size_t last = i - 1; + size_t tmp = indices[last]; + indices[last] = indices[j]; + indices[j] = tmp; + } + return 1; +} diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 6840e10..38ae4a9 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -26,6 +26,7 @@ nir_add_test(test_rng) nir_add_test(test_encode) nir_add_test(test_data) nir_add_test(test_checkpoint) +nir_add_test(test_minibatch) # 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_minibatch.c b/tests/test_minibatch.c new file mode 100644 index 0000000..dd36557 --- /dev/null +++ b/tests/test_minibatch.c @@ -0,0 +1,299 @@ +/* Mini-batch training: accumulation over a caller loop equals the + * sum of per-sample gradients, the batch mean matches the + * single-sample losses, and epoch shuffling replays from its + * seed. References are float64 closed forms; tolerances separate + * float32 rounding from formula errors. + */ +#include +#include +#include + +#include "nir_c_runtime/data.h" +#include "nir_c_runtime/minibatch.h" +#include "nir_c_runtime/ops.h" +#include "nir_c_runtime/rng.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 test_zero_accumulate_scale(void) { + /* Dyadic fractions, so the sums below are exact in float32. */ + float g1[] = {1.5f, -2.25f, 3.5f}; + float g2[] = {0.5f, 1.25f, -1.5f}; + float g3[] = {-1.0f, 0.5f, -1.0f}; + float accum[3]; + nir_minibatch_zero(accum, 3); + CHECK(accum[0] == 0.0f && accum[1] == 0.0f && accum[2] == 0.0f); + CHECK(nir_minibatch_accumulate(accum, g1, 3)); + CHECK(nir_minibatch_accumulate(accum, g2, 3)); + CHECK(nir_minibatch_accumulate(accum, g3, 3)); + CHECK(accum[0] == 1.0f && accum[1] == -0.5f && accum[2] == 1.0f); + CHECK(nir_minibatch_scale(accum, 3, 1.0f / 3.0f)); + CHECK_NEAR(accum[0], 1.0 / 3.0, 1e-7); + CHECK_NEAR(accum[1], -1.0 / 6.0, 1e-7); + CHECK_NEAR(accum[2], 1.0 / 3.0, 1e-7); +} + +/* A batch of B=4 affine samples: the helper path (zero, then one + * accumulate per sample) must match the sum of the per-sample + * gradients, checked against a float64 reference and against the + * op's own repeated-call accumulation. */ +static void test_accumulate_matches_affine_loop(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 xs[4][2] = { + {0.5f, -0.25f}, {1.0f, 0.5f}, {-1.5f, 2.0f}, {0.125f, -1.0f}, + }; + float gs[4][2] = { + {1.0f, -0.5f}, {0.25f, 0.75f}, {-2.0f, 1.0f}, {0.5f, 0.5f}, + }; + size_t mshape[2] = {2, 2}; + size_t vshape[1] = {2}; + nir_tensor_t sample_dw = nir_tensor_alloc(mshape, 2); + nir_tensor_t sample_db = nir_tensor_alloc(vshape, 1); + nir_tensor_t direct_dw = nir_tensor_alloc(mshape, 2); + nir_tensor_t direct_db = nir_tensor_alloc(vshape, 1); + nir_tensor_t grad_in = nir_tensor_alloc(vshape, 1); + CHECK(sample_dw.data != NULL && sample_db.data != NULL); + CHECK(direct_dw.data != NULL && direct_db.data != NULL); + CHECK(grad_in.data != NULL); + float accum_w[4]; + float accum_b[2]; + double ref_w[4] = {0.0, 0.0, 0.0, 0.0}; + double ref_b[2] = {0.0, 0.0}; + nir_minibatch_zero(accum_w, 4); + nir_minibatch_zero(accum_b, 2); + for (size_t k = 0; k < 4; k++) { + nir_tensor_t input = make_vector(xs[k], 2); + nir_tensor_t gout = make_vector(gs[k], 2); + memset(sample_dw.data, 0, 4 * sizeof(float)); + memset(sample_db.data, 0, 2 * sizeof(float)); + CHECK(nir_op_affine_backward( + ¶ms, &input, &gout, &sample_dw, &sample_db, + &grad_in + )); + /* The direct path accumulates across calls untouched. */ + CHECK(nir_op_affine_backward( + ¶ms, &input, &gout, &direct_dw, &direct_db, + &grad_in + )); + CHECK(nir_minibatch_accumulate( + accum_w, sample_dw.data, 4 + )); + CHECK(nir_minibatch_accumulate( + accum_b, sample_db.data, 2 + )); + for (size_t i = 0; i < 4; i++) { + ref_w[i] += (double)sample_dw.data[i]; + } + for (size_t i = 0; i < 2; i++) { + ref_b[i] += (double)sample_db.data[i]; + } + nir_tensor_free(&input); + nir_tensor_free(&gout); + } + for (size_t i = 0; i < 4; i++) { + CHECK_NEAR(accum_w[i], ref_w[i], 1e-6); + CHECK(accum_w[i] == direct_dw.data[i]); + } + for (size_t i = 0; i < 2; i++) { + CHECK_NEAR(accum_b[i], ref_b[i], 1e-6); + CHECK(accum_b[i] == direct_db.data[i]); + } + CHECK(nir_minibatch_scale(accum_w, 4, 1.0f / 4.0f)); + CHECK(nir_minibatch_scale(accum_b, 2, 1.0f / 4.0f)); + for (size_t i = 0; i < 4; i++) { + CHECK_NEAR(accum_w[i], ref_w[i] / 4.0, 1e-6); + } + for (size_t i = 0; i < 2; i++) { + CHECK_NEAR(accum_b[i], ref_b[i] / 4.0, 1e-6); + } + nir_tensor_free(&sample_dw); + nir_tensor_free(&sample_db); + nir_tensor_free(&direct_dw); + nir_tensor_free(&direct_db); + nir_tensor_free(&grad_in); +} + +/* A batch of B=4 cross-entropy samples: the mean loss and the + * mean gradient match the per-sample single-sample CE values. */ +static void test_mean_loss_matches_ce(void) { + float ls[4][3] = { + {1.0f, 2.0f, 3.0f}, + {0.5f, -1.0f, 0.25f}, + {-2.0f, 0.0f, 1.0f}, + {3.0f, 3.0f, 3.0f}, + }; + size_t targets[4] = {2, 0, 1, 2}; + size_t shape[1] = {3}; + nir_tensor_t grad = nir_tensor_alloc(shape, 1); + CHECK(grad.data != NULL); + float losses[4]; + float accum[3]; + double ref_loss = 0.0; + double ref_grad[3] = {0.0, 0.0, 0.0}; + nir_minibatch_zero(accum, 3); + for (size_t k = 0; k < 4; k++) { + nir_tensor_t logits = make_vector(ls[k], 3); + CHECK(nir_cross_entropy( + &logits, targets[k], &losses[k], &grad + )); + CHECK(nir_minibatch_accumulate(accum, grad.data, 3)); + ref_loss += (double)losses[k]; + for (size_t i = 0; i < 3; i++) { + ref_grad[i] += (double)grad.data[i]; + } + nir_tensor_free(&logits); + } + float mean = 0.0f; + CHECK(nir_minibatch_mean_loss(losses, 4, &mean)); + CHECK_NEAR(mean, ref_loss / 4.0, 1e-7); + CHECK(nir_minibatch_scale(accum, 3, 1.0f / 4.0f)); + for (size_t i = 0; i < 3; i++) { + CHECK_NEAR(accum[i], ref_grad[i] / 4.0, 1e-6); + } + nir_tensor_free(&grad); +} + +static void check_permutation( + const size_t *indices, size_t n +) { + int *seen = (int *)calloc(n == 0 ? 1 : n, sizeof(int)); + CHECK(seen != NULL); + for (size_t i = 0; i < n; i++) { + CHECK(indices[i] < n); + CHECK(!seen[indices[i]]); + seen[indices[i]] = 1; + } + free(seen); +} + +static void test_shuffle_reproducible(void) { + size_t a[16]; + size_t b[16]; + size_t c[16]; + nir_rng_t rng; + nir_sampler_arange(a, 16); + nir_rng_seed(&rng, 42); + CHECK(nir_sampler_shuffle(&rng, a, 16)); + check_permutation(a, 16); + nir_sampler_arange(b, 16); + nir_rng_seed(&rng, 42); + CHECK(nir_sampler_shuffle(&rng, b, 16)); + CHECK(memcmp(a, b, sizeof(a)) == 0); + nir_sampler_arange(c, 16); + nir_rng_seed(&rng, 43); + CHECK(nir_sampler_shuffle(&rng, c, 16)); + check_permutation(c, 16); + CHECK(memcmp(a, c, sizeof(a)) != 0); +} + +/* Epoch shuffling over the committed tiny-MNIST fixture: the + * order replays from its seed and indexes the real samples. */ +static void test_shuffle_fixture(void) { + char *error = NULL; + uint8_t *pixels = NULL; + uint8_t *labels = NULL; + size_t count = 0; + size_t rows = 0; + size_t cols = 0; + size_t nlabels = 0; + CHECK(nir_idx_read_images( + "fixtures/mnist_tiny_images.idx", &pixels, &count, &rows, + &cols, &error + )); + CHECK(error == NULL && count == 3); + CHECK(nir_idx_read_labels( + "fixtures/mnist_tiny_labels.idx", &labels, &nlabels, &error + )); + CHECK(error == NULL && nlabels == count); + size_t first[3]; + size_t second[3]; + nir_rng_t rng; + nir_sampler_arange(first, count); + nir_rng_seed(&rng, 7); + CHECK(nir_sampler_shuffle(&rng, first, count)); + check_permutation(first, count); + /* Pinned order for seed 7 (Fisher-Yates over splitmix64). */ + CHECK(first[0] == 1 && first[1] == 2 && first[2] == 0); + nir_sampler_arange(second, count); + nir_rng_seed(&rng, 7); + CHECK(nir_sampler_shuffle(&rng, second, count)); + CHECK(memcmp(first, second, sizeof(first)) == 0); + /* The order indexes real samples: labels 7,3,9 read as 3,9,7. */ + CHECK(labels[0] == 7 && labels[1] == 3 && labels[2] == 9); + CHECK(labels[first[0]] == 3 && labels[first[1]] == 9 && + labels[first[2]] == 7); + nir_tensor_t frame = nir_data_frame( + pixels, count, rows, cols, first[0] + ); + CHECK(frame.data != NULL && frame.size == rows * cols); + nir_tensor_free(&frame); + free(pixels); + free(labels); +} + +static void test_shuffle_edges(void) { + nir_rng_t rng; + nir_rng_seed(&rng, 1); + CHECK(nir_sampler_shuffle(&rng, NULL, 0)); + CHECK(!nir_sampler_shuffle(NULL, NULL, 0)); + size_t one[1] = {0}; + CHECK(nir_sampler_shuffle(&rng, one, 1)); + CHECK(one[0] == 0); + size_t two[2] = {0, 1}; + CHECK(!nir_sampler_shuffle(NULL, two, 2)); + CHECK(!nir_sampler_shuffle(&rng, NULL, 2)); +} + +static void test_errors(void) { + float accum[3] = {1.0f, 2.0f, 3.0f}; + float grad[3] = {0.5f, 0.5f, 0.5f}; + float bad[3] = {0.5f, (float)NAN, 0.5f}; + float inf[3] = {0.5f, (float)INFINITY, 0.5f}; + float before[3]; + memcpy(before, accum, sizeof(before)); + CHECK(!nir_minibatch_accumulate(NULL, grad, 3)); + CHECK(!nir_minibatch_accumulate(accum, NULL, 3)); + CHECK(!nir_minibatch_accumulate(accum, grad, 0)); + CHECK(!nir_minibatch_accumulate(accum, bad, 3)); + CHECK(!nir_minibatch_accumulate(accum, inf, 3)); + CHECK(!nir_minibatch_accumulate(bad, grad, 3)); + CHECK(memcmp(accum, before, sizeof(before)) == 0); + CHECK(!nir_minibatch_scale(NULL, 3, 0.5f)); + CHECK(!nir_minibatch_scale(accum, 0, 0.5f)); + CHECK(!nir_minibatch_scale(accum, 3, (float)NAN)); + CHECK(!nir_minibatch_scale(accum, 3, (float)INFINITY)); + CHECK(!nir_minibatch_scale(bad, 3, 0.5f)); + CHECK(memcmp(accum, before, sizeof(before)) == 0); + CHECK(nir_minibatch_scale(accum, 3, 0.0f)); + CHECK(accum[0] == 0.0f && accum[1] == 0.0f && accum[2] == 0.0f); + float losses[2] = {1.0f, 3.0f}; + float mean = 0.0f; + CHECK(!nir_minibatch_mean_loss(NULL, 2, &mean)); + CHECK(!nir_minibatch_mean_loss(losses, 0, &mean)); + CHECK(!nir_minibatch_mean_loss(losses, 2, NULL)); + CHECK(!nir_minibatch_mean_loss(bad, 2, &mean)); + CHECK(!nir_minibatch_mean_loss(inf, 2, &mean)); + CHECK(nir_minibatch_mean_loss(losses, 2, &mean)); + CHECK(mean == 2.0f); +} + +int main(void) { + test_zero_accumulate_scale(); + test_accumulate_matches_affine_loop(); + test_mean_loss_matches_ce(); + test_shuffle_reproducible(); + test_shuffle_fixture(); + test_shuffle_edges(); + test_errors(); + return 0; +}