1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
|
#include "utils.h"
#include "dynmlp.h"
#include "ode_solver.h"
#include "adjoint.h"
#include "adam.h"
#include "train.h"
#include "spiral.h"
#include "tests.h"
#include "test_cnf.h"
#include "cnf_train.h"
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
int main(void) {
RNG r = rng_init(42);
/* --- sanity checks --- */
test_ode_solver();
test_dynmlp_gradients(&r);
test_adjoint_gradients(&r);
test_multi_obs_adjoint(&r);
test_training(&r);
printf("\n--- CNF tests ---\n");
test_cnf_trace(&r);
test_cnf_invertibility(&r);
test_cnf_gradients(&r);
test_cnf_training(&r);
cnf_train_demo(&r);
printf("\n--- Training demo (spiral) ---\n\n");
r = rng_init((uint64_t)time(NULL));
const double t0 = 0.0, t1 = 1.5;
const double noise_std = 0.1;
const double atol_train = 1e-3, rtol_train = 1e-3;
const double atol_eval = 1e-5, rtol_eval = 1e-5;
const int TRAIN_N = 200, TEST_N = 50;
const int ITERS = 500, BATCH = 16, LOG_EVERY = 25;
Dataset train_ds = generate_spiral_dataset(TRAIN_N, t0, t1, noise_std, &r);
Dataset test_ds = generate_spiral_dataset(TEST_N, t0, t1, noise_std, &r);
const int D = 2, H = 32;
DynMLP net;
int nparams = dynmlp_nparams(D, H);
double *theta = vec_alloc(nparams);
dynmlp_init(&net, D, H, theta, &r);
Adam adam = adam_init(nparams, 1e-3, 0.9, 0.999, 1e-8);
printf("%-6s %-12s %-12s %-10s %-10s",
"Iter", "Train Loss", "Test Loss", "Fwd NFE", "Bwd NFE");
printf("\n--------------------------------------------------\n");
const double **batch_z0 = (const double **)xmalloc(BATCH * sizeof(double *));
const double **batch_tgt = (const double **)xmalloc(BATCH * sizeof(double *));
for (int iter = 1; iter <= ITERS; iter++) {
for (int b = 0; b < BATCH; b++) {
int idx = (int)(rng_next(&r) % (uint64_t)TRAIN_N);
batch_z0[b] = train_ds.z0[idx];
batch_tgt[b] = train_ds.target[idx];
}
TrainStepResult res = train_step(&net, theta, batch_z0, batch_tgt,
t0, t1, BATCH, &adam,
atol_train, rtol_train, 10);
if (iter % LOG_EVERY == 0) {
double test_loss = evaluate(&net, theta, &test_ds,
t0, t1, atol_eval, rtol_eval);
printf("%-6d %-12.6f %-12.6f %-10d %-10d\n",
iter, res.loss, test_loss,
res.nfe_fwd, res.nfe_bwd);
fflush(stdout);
}
}
free(batch_z0);
free(batch_tgt);
printf("\n--------------------------------------------------\n");
double final_test_loss = evaluate(&net, theta, &test_ds,
t0, t1, atol_eval, rtol_eval);
printf("Final test loss : %.6f\n", final_test_loss);
printf("Total parameters: %d\n", nparams);
printf("\nSample predictions:\n");
Workspace ws = workspace_alloc(D, H, nparams);
AdjointCtx ac = { net, theta, D, nparams, &ws };
for (int s = 0; s < 5; s++) {
int idx = (int)(rng_next(&r) % (uint64_t)TEST_N);
ODEResult fwd = ode_solve(neural_ode_rhs,
test_ds.z0[idx], t0, t1,
NULL, D, atol_eval, rtol_eval, &ac);
printf(" z0=(%.4f, %.4f) predicted=(%.4f, %.4f) target=(%.4f, %.4f)\n",
test_ds.z0[idx][0], test_ds.z0[idx][1],
fwd.y[0], fwd.y[1],
test_ds.target[idx][0], test_ds.target[idx][1]);
free(fwd.y);
}
workspace_free(&ws);
free(theta);
adam_free(&adam);
dataset_free(&train_ds);
dataset_free(&test_ds);
return 0;
}
|