This repository contains the pipeline and data sets supporting the results of the following article:
Fourment M, Swanepoel CJ, Galloway JG, Ji X, Gangavarapu K, Suchard MA, Matsen IV FA. Automatic differentiation is no panacea for phylogenetic gradient computation. arXiv:2211.02168
This benchmark compares the efficiency (memory usage and speed) of several gradient implementations of phylogenetic models (e.g., tree likelikelihood and coalescent model). The goal of this study is to compare the efficiency of automatic differentiation (AD) and analytic gradient. The pipeline reuses parts of the treetime validation workflow.
| Program | Language | Framework | Gradient | BITO support |
|---|---|---|---|---|
| physher | C | analytic | ||
| phylostan | Stan | Stan | AD | |
| phylojax | python | JAX | AD | |
| torchtree | python | PyTorch | AD | ✅ |
| treeflow | python | TensorFlow | AD |
The gradient of the tree likelihood is optionaly computed by BITO, an efficient C++ library that analytically calculate the gradient using the BEAGLE library. torchtree uses the torchtree-bito plugin to access BITO.
You will need to install nextflow and docker to run this benchmark. Docker is not required but it is highly recommended to use it due to the numerous dependencies.
git clone 4ment/autodiff-experiments.git
git submodule update --init --recursive
nextflow run 4ment/autodiff-experiments -profile docker -with-trace
Since the pipeline will take weeks to run to completion one should use a high performance computer. Examples of configuration files for pbspro and slurm can be found in the configs folder.
Before generating the figures, we need to extract memory usage information from the trace.txt file and work directory:
python scripts/parse-trace.py work/ trace.txt > results/trace.csv
Generate figures in a single pdf:
Rscript -e 'rmarkdown::render("plot.Rmd")'
For reproducbility, we provide below the version or commit hash of each library/program used in the benchmark.
| Library | Version |
|---|---|
| jax | 0.2.24 |
| jaxlib | 0.3.7 |
| numpy | 1.22 |
| pystan | 2.19.1.1 |
| tensorflow | 2.10.0 |
| tensorflow-probability | 0.18.0 |
| pytorch | 1.12.1 |
| Program | Version/hash |
|---|---|
| bito | cc0806abcd0b9f2fab604e800c674c9a5c5afebe |
| phylojax | a1612cae36292af76e8d24cc40d6544162c987aa |
| phylostan | 1.0.5 |
| physher | b19ff2f9422f29ba1ab31306a3fe29ab6a6f607b |
| torchtree | f3831650a807e74cc2e9478009e57a41f47bed8d |
| torchtree-bito | e2a95cefb13968f95f6e5520bd0a52d726ee7fc9 |
| treeflow | e3414dcc9e764d06abc3e19c1d0f55110499e2ea |