tfNeuralODE-Adjoint

In this example, we’re going to train a simple network to learn a new spiral trajectory using the adjoint method to train. We start by loading all of the libraries and setting our initial conditions, along with a few constants.

library(reticulate)
library(tensorflow)
library(tfNeuralODE)
library(keras)
library(deSolve)

# iterations, time span (layers)
niters = 25
t = seq(0, 25, by = 25/100)
# initial conditions
h0 = tf$cast(t(c(1., 0.)), dtype= tf$float32)
W = tf$cast(rbind(c(-0.1, 1.0), c(-0.2, -0.1)), dtype = tf$float32)
h0_var = tf$Variable(h0, name = "")
hN_target = tf$cast(t(c(0., 0.5)), dtype = tf$float32)

We solve for the initial trajectory.

trueODEfunc<- function(du, u, p, t){
  true_A = rbind(c(-0.1, 1.0), c(-0.2, -0.1))
  du <- (u) %*% true_A
  return(list(du))
}

# solved ode output
init_path <- lsode(func = trueODEfunc, y = c(1., 0.), times = t)

Now we instantiate a very simple ODE model following that initial trajectory and an optimizer to train the ODE model with.

# ODE Model

optimizer = tf$keras$optimizers$legacy$SGD(learning_rate=1e-2, momentum=0.95)

OdeModel(keras$Model) %py_class% {
  call <- function(inputs) {
    tf$matmul(inputs, W)
  }
}
model<- OdeModel()

Now we train the model, using 25 iterations. Each 5 iterations, the plot of the ODE will be produced.

for(i in 1:niters){
  print(paste("Iteration", i, "out of", niters, "iterations."))
  with(tf$GradientTape() %as% tape, {
    pred = forward(model, inputs = h0_var, tsteps = t)
    tape$watch(pred)
    loss = tf$reduce_sum((hN_target - pred) ^ 2)
  })
  #print(paste("loss:", as.numeric(loss)))
  dLoss = tape$gradient(loss, pred)
  dfdh0 = backward(model, t, pred, output_gradients = dLoss)[[2]]
  optimizer$apply_gradients(list(c(dfdh0, h0_var)))

  # graphing the Neural ODE
  if(i %% 5 == 0 || i == 1){
    pred_y = forward(model = model, inputs = tf$cast((as.matrix(h0_var)), dtype = tf$float32),
                     tsteps = t, return_states = TRUE)
    pred_y_c<- k_concatenate(pred_y[[2]], 1)
    p_m<- as.matrix(pred_y_c)
    plot(p_m,
         main = paste("Iteration", i), type = "l", col = "red",
         xlim = c(-1,2), ylim = c(-1,2))
    lines(init_path[,2], init_path[,3], col = "blue")
  }
}
#> [1] "Iteration 1 out of 25 iterations."

plot of iteration 1

#> [1] "Iteration 2 out of 25 iterations."
#> [1] "Iteration 3 out of 25 iterations."
#> [1] "Iteration 4 out of 25 iterations."
#> [1] "Iteration 5 out of 25 iterations."

plot of iteration 5

#> [1] "Iteration 6 out of 25 iterations."
#> [1] "Iteration 7 out of 25 iterations."
#> [1] "Iteration 8 out of 25 iterations."
#> [1] "Iteration 9 out of 25 iterations."
#> [1] "Iteration 10 out of 25 iterations."

plot of iteration 10

#> [1] "Iteration 11 out of 25 iterations."
#> [1] "Iteration 12 out of 25 iterations."
#> [1] "Iteration 13 out of 25 iterations."
#> [1] "Iteration 14 out of 25 iterations."
#> [1] "Iteration 15 out of 25 iterations."

plot of iteration 15

#> [1] "Iteration 16 out of 25 iterations."
#> [1] "Iteration 17 out of 25 iterations."
#> [1] "Iteration 18 out of 25 iterations."
#> [1] "Iteration 19 out of 25 iterations."
#> [1] "Iteration 20 out of 25 iterations."

plot of iteration 20

#> [1] "Iteration 21 out of 25 iterations."
#> [1] "Iteration 22 out of 25 iterations."
#> [1] "Iteration 23 out of 25 iterations."
#> [1] "Iteration 24 out of 25 iterations."
#> [1] "Iteration 25 out of 25 iterations."

plot of iteration 25