Need help with easy-neural-ode?
Click the “chat” button below for chat support from the developer who created it, or find similar developers for support.

About the developer

jacobjinkelly
136 Stars 12 Forks MIT License 74 Commits 1 Opened issues

Description

Code for the paper "Learning Differential Equations that are Easy to Solve"

Services available

!
?

Need anything else?

Contributors list

# 120,627
Jupyter...
C++
React
TeX
71 commits

Learning Differential Equations that are Easy to Solve

Code for the paper:

Jacob Kelly, Jesse Bettencourt, Matthew James Johnson, David Duvenaud. "Learning Differential Equations that are Easy to Solve." Neural Information Processing Systems (2020). [arxiv] [bibtex]

*Equal Contribution

Includes JAX implementations of the following models: - Neural ODEs for classification - Latent ODEs for time series - FFJORD for density estimation

Includes JAX implementations of the following adaptive-stepping numerical solvers: - Heun-Euler

heun
(2nd order) - Fehlberg (RK1(2))
fehlberg
(2nd order) - Bogacki-Shampine
bosh
(3rd order) - Cash-Karp
cash_karp
(4th order) - Fehlberg
rk_fehlberg
(4th order) - Owrenzen
owrenzen
(4th order) - Dormand-Prince
dopri
(5th order) - Owrenzen
owrenzen5
(5th order) - Tanyam
tanyam
(7th order) - Adams
adams
(adaptive order) - RK4
rk4
(4th order, fixed step-size)

Requirements

JAX

Follow installation instructions here.

Haiku

Follow installation instructions here.

Tensorflow Datasets

For using the MNIST dataset, follow installation instructions here.

Usage

Different scripts are provided for each task and dataset.

MNIST Classification

python mnist.py --reg r3 --lam 6e-5

Latent ODEs

python latent_ode.py --reg r3 --lam 1e-2

FFJORD (Tabular)

python ffjord_tabular.py --reg r2 --lam 1e-2

FFJORD (MNIST)

python ffjord_mnist.py --reg r2 --lam 3e-4

Datasets

MNIST

tensorflow-datasets
(instructions for installing above) will download the data when called from the training script.

Physionet

The file

physionet_data.py
, adapted from Latent ODEs for Irregularly-Sampled Time Series will download and process the data when called from the training script. A preprocessed version is available in releases.

Tabular (FFJORD)

Data must be downloaded following instructions from gpapamak/maf and placed in

data/
. Only
MINIBOONE
is needed for experiments in the paper.

Code in

datasets/
, adapted from Free-form Jacobian of Reversible Dynamics (FFJORD), will create an interface for the
MINIBOONE
dataset once it's downloaded. It is called from the training script.

Acknowledgements

Code in

lib
is modified from google/jax under the license.

Several numerical solvers were adapted from torchdiffeq and DifferentialEquations.jl.

BibTeX

@inproceedings{kelly2020easynode,
  title={Learning Differential Equations that are Easy to Solve},
  author={Kelly, Jacob and Bettencourt, Jesse and Johnson, Matthew James and Duvenaud, David},
  booktitle={Neural Information Processing Systems},
  year={2020},
  url={https://arxiv.org/abs/2007.04504}
}

We use cookies. If you continue to browse the site, you agree to the use of cookies. For more information on our use of cookies please see our Privacy Policy.