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
|
#include "train.h"
#include "adjoint.h"
#include "utils.h"
#include <stdlib.h>
static double train_one(DynNet *net, const double *theta,
const double *z0, double t0, double t1,
const double *target, double *grad_accum,
double atol, double rtol, int num_checkpoints,
int *nfe_fwd, int *nfe_bwd) {
NeuralODEOutput out = neural_ode_forward_backward(net, theta, z0, t0, t1,
target, atol, rtol, num_checkpoints);
double loss = 0.0;
int D = net->D;
for (int i = 0; i < D; i++) {
double d = out.z1[i] - target[i];
loss += 0.5 * d * d;
}
for (int i = 0; i < net->total_params; i++) grad_accum[i] += out.dL_dtheta[i];
*nfe_fwd += out.nfe_forward;
*nfe_bwd += out.nfe_backward;
free(out.z1);
free(out.dL_dz0);
free(out.dL_dtheta);
return loss;
}
TrainStepResult train_step(DynNet *net, double *theta,
const double **z0s, const double **targets,
double t0, double t1, int batch_size,
Adam *adam, double atol, double rtol, int num_checkpoints) {
int nparams = net->total_params;
double *grad_accum = vec_zeros(nparams);
TrainStepResult res = { 0.0, 0, 0 };
for (int b = 0; b < batch_size; b++) {
res.loss += train_one(net, theta, z0s[b], t0, t1, targets[b],
grad_accum, atol, rtol, num_checkpoints,
&res.nfe_fwd, &res.nfe_bwd);
}
res.loss /= (double)batch_size;
for (int i = 0; i < nparams; i++) grad_accum[i] /= (double)batch_size;
adam_update(adam, theta, grad_accum);
free(grad_accum);
return res;
}
|