DynGFN: Bayesian Dynamic Causal Discovery using Generative Flow Networks

Description

Generative Flow Networks (GFlowNets) are an approach for learning generative models over discrete spaces. GFlowNets learn a stochastic policy $P_F(\tau)$ to sequentially sample an object $\mathbf{x}$ (e.g. a graph) from a discrete space $\mathcal{X}$. We propose a general framework for cyclic causal discovery of dynamic systems and use GFlowNets to learn dynamic causal structure. See this http link for the paper.

DynGFN is method for learning cyclic causal structure of dynamic systems from observational dynamic data. DynGFN consists of three components: (1) A graph sampler that samples graphical structure encoding the causal dependencies between variables, (2) A class of structural equation modules that models the functional relationships between the observed variables, indexed by parameters $\theta$, and (3) A hyper-network architecture that outputs the parameters $\theta$ of the structural equations as a function of the sampled graph, i.e. $\theta = h_{\phi}(G)$.

If you find this code useful in your research, please cite the following paper (expand for BibTeX):

L. Atanackovic*, A. Tong*, B. Wang, L.J. Lee, Y. Bengio, J. Hartford. DynGFN: Towards Bayesian Inference of Gene Regulatory Networks with GFlowNets, NeurIPS, 2023.

@article{atanackovic2023dyngfn,
  title={DynGFN: Towards Bayesian Inference of Gene Regulatory Networks with GFlowNets},
  author={Atanackovic, Lazar and Tong, Alexander and Wang, Bo and Lee, Leo J and Bengio, Yoshua and Hartford, Jason S},
  journal={Advances in Neural Information Processing Systems},
  volume={36},
  year={2023}
}

How to run

Install dependencies

# clone project
git clone https://github.com/lazaratan/dyn-gfn.git
cd dyn-gfn

# [OPTIONAL] create conda environment
conda create -n myenv python=3.8
conda activate myenv

# install pytorch according to instructions
# https://pytorch.org/get-started/

# install requirements
pip install -r requirements.txt

Train model with default configuration

# train on CPU
python train.py

# train on GPU
python train.py trainer=gpu

Train model with chosen experiment configuration from configs/experiment/

python train.py experiment=experiment_name.yaml

You can override any parameter from command line like this

python train.py trainer.max_epochs=20 datamodule.batch_size=64

Project Structure

The directory structure of new project looks like this:

├── configs                   <- Hydra configuration files
│   ├── callbacks                <- Callbacks configs
│   ├── datamodule               <- Datamodule configs
│   ├── debug                    <- Debugging configs
│   ├── experiment               <- Experiment configs
│   ├── hparams_search           <- Hyperparameter search configs
│   ├── local                    <- Local configs
│   ├── log_dir                  <- Logging directory configs
│   ├── logger                   <- Logger configs
│   ├── model                    <- Model configs
│   ├── trainer                  <- Trainer configs
│   │
│   ├── test.yaml             <- Main config for testing
│   └── train.yaml            <- Main config for training
│
├── data                   <- Project data
│
├── logs                   <- Logs generated by Hydra and PyTorch Lightning loggers
│
├── notebooks              <- Jupyter notebooks. Naming convention is a number (for ordering),
│                             the creator's initials, and a short `-` delimited description,
│                             e.g. `1.0-jqp-initial-data-exploration.ipynb`.
│
├── scripts                <- Shell scripts
│
├── src                    <- Source code
│   ├── datamodules              <- Lightning datamodules
│   ├── models                   <- Lightning models
│   ├── utils                    <- Utility scripts
│   ├── vendor                   <- Third party code that cannot be installed using PIP/Conda
│   │
│   ├── testing_pipeline.py
│   └── training_pipeline.py
│
├── tests                  <- Tests of any kind
│   ├── helpers                  <- A couple of testing utilities
│   ├── shell                    <- Shell/command based tests
│   └── unit                     <- Unit tests
│
├── test.py               <- Run testing
├── train.py              <- Run training
│
├── .env.example              <- Template of the file for storing private environment variables
├── .gitignore                <- List of files/folders ignored by git
├── .pre-commit-config.yaml   <- Configuration of pre-commit hooks for code formatting
├── requirements.txt          <- File for installing python dependencies
├── setup.cfg                 <- Configuration of linters and pytest
└── README.md

⚡  Your Superpowers

Override any config parameter from command line

Hydra allows you to easily overwrite any parameter defined in your config.

python train.py trainer.max_epochs=20 model.lr=1e-4

You can also add new parameters with + sign.

python train.py +model.new_param="uwu"

Train on CPU, GPU, multi-GPU and TPU

PyTorch Lightning makes it easy to train your models on different hardware.

# train on CPU
python train.py trainer.gpus=0

# train on 1 GPU
python train.py trainer.gpus=1

# train on TPU
python train.py +trainer.tpu_cores=8

# train with DDP (Distributed Data Parallel) (4 GPUs)
python train.py trainer.gpus=4 +trainer.strategy=ddp

# train with DDP (Distributed Data Parallel) (8 GPUs, 2 nodes)
python train.py trainer.gpus=4 +trainer.num_nodes=2 +trainer.strategy=ddp

Train with mixed precision

# train with pytorch native automatic mixed precision (AMP)
python train.py trainer.gpus=1 +trainer.precision=16

Train model with any logger available in PyTorch Lightning, like Weights&Biases or Tensorboard

PyTorch Lightning provides convenient integrations with most popular logging frameworks, like Tensorboard, Neptune or simple csv files. Read more here. Using wandb requires you to setup account first. After that just complete the config as below.
> Click here to see example wandb dashboard generated with this template.

# set project and entity names in `configs/logger/wandb`
wandb:
  project: "your_project_name"
  entity: "your_wandb_team_name"

# train model with Weights&Biases (link to wandb dashboard should appear in the terminal)
python train.py logger=wandb

Train model with chosen experiment config

Experiment configurations are placed in configs/experiment/.

python train.py experiment=example

Attach some callbacks to run

Callbacks can be used for things such as as model checkpointing, early stopping and many more.
Callbacks configurations are placed in configs/callbacks/.

python train.py callbacks=default

Use different tricks available in Pytorch Lightning

PyTorch Lightning provides about 40+ useful trainer flags.

# gradient clipping may be enabled to avoid exploding gradients
python train.py +trainer.gradient_clip_val=0.5

# stochastic weight averaging can make your models generalize better
python train.py +trainer.stochastic_weight_avg=true

# run validation loop 4 times during a training epoch
python train.py +trainer.val_check_interval=0.25

# accumulate gradients
python train.py +trainer.accumulate_grad_batches=10

# terminate training after 12 hours
python train.py +trainer.max_time="00:12:00:00"

Easily debug

Visit configs/debug/ for different debugging configs.

# runs 1 epoch in default debugging mode
# changes logging directory to `logs/debugs/...`
# sets level of all command line loggers to 'DEBUG'
# enables extra trainer flags like tracking gradient norm
# enforces debug-friendly configuration
python train.py debug=default

# runs test epoch without training
python train.py debug=test_only

# run 1 train, val and test loop, using only 1 batch
python train.py +trainer.fast_dev_run=true

# raise exception if there are any numerical anomalies in tensors, like NaN or +/-inf
python train.py +trainer.detect_anomaly=true

# print execution time profiling after training ends
python train.py +trainer.profiler="simple"

# try overfitting to 1 batch
python train.py +trainer.overfit_batches=1 trainer.max_epochs=20

# use only 20% of the data
python train.py +trainer.limit_train_batches=0.2 \
+trainer.limit_val_batches=0.2 +trainer.limit_test_batches=0.2

# log second gradient norm of the model
python train.py +trainer.track_grad_norm=2

Resume training from checkpoint

Checkpoint can be either path or URL.

python train.py trainer.resume_from_checkpoint="/path/to/ckpt/name.ckpt"

⚠️ Currently loading ckpt in Lightning doesn't resume logger experiment, but it will be supported in future Lightning release.

Execute evaluation for a given checkpoint

Checkpoint can be either path or URL.

python test.py ckpt_path="/path/to/ckpt/name.ckpt"

Create a sweep over hyperparameters

# this will run 6 experiments one after the other,
# each with different combination of batch_size and learning rate
python train.py -m datamodule.batch_size=32,64,128 model.lr=0.001,0.0005

⚠️ This sweep is not failure resistant (if one job crashes than the whole sweep crashes).

Create a sweep over hyperparameters with Optuna

Using Optuna Sweeper plugin doesn't require you to code any boilerplate into your pipeline, everything is defined in a single config file!

# this will run hyperparameter search defined in `configs/hparams_search/mnist_optuna.yaml`
# over chosen experiment config
python train.py -m hparams_search=mnist_optuna experiment=example_simple

⚠️ Currently this sweep is not failure resistant (if one job crashes than the whole sweep crashes). Might be supported in future Hydra release.

Execute all experiments from folder

Hydra provides special syntax for controlling behavior of multiruns. Learn more here. The command below executes all experiments from folder configs/experiment/.

python train.py -m 'experiment=glob(*)'

Execute sweep on a remote AWS cluster

This should be achievable with simple config using Ray AWS launcher for Hydra. Example is not yet implemented in this template.

Use Hydra tab completion

Hydra allows you to autocomplete config argument overrides in shell as you write them, by pressing tab key. Learn more here.

Apply pre-commit hooks

Apply pre-commit hooks to automatically format your code and configs, perform code analysis and remove output from jupyter notebooks. See # Best Practices for more.

pre-commit run -a

❤️  Contributions

Have a question? Found a bug? Missing a specific feature? Feel free to file a new issue, discussion or PR with respective title and description.

Before making an issue, please verify that:

• The problem still exists on the current main branch.
• Your python dependencies are updated to recent versions.

Suggestions for improvements are always welcome!

License

This project is licensed under the MIT License.

MIT License

Copyright (c) 2023 Lazar Atanackovic

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.