Hello,
I am training a conditional Universal Differential Equation in Julia. This is a regular UDE with an additional trainable parameter as input that is specific for each particle (individual) in the dataset.
The model has one ODE state, a small shared neural network, and one particle-specific conditional parameter for each of 600 simulated particle trajectories. Training currently takes several hours and training time increases with the number of particles.
The pipeline consists of:
- Jointly optimising the neural-network parameters and 600 conditional parameters using ADAM followed by L-BFGS (this is the bottle-neck)
- Reestimating each conditional parameter separately while keeping the neural network parameters fixed
- Parallelising the multistart runs with
pmap
I would like to know whether this workflow can be made substantially faster without affecting the outcomes. I would appreciate any feedback that can help speed up this code (hopefully) by orders of magnitude.
Let me know if anything is unclear or required as I’m new to the Julia Discourse.
Thank you for your time!
Kind regards,
NCCP
dBPM.jl (13.2 KB)
inputs.jl (7.3 KB)
ude.jl (18.1 KB)