Implicit Deep Adaptive Design (iDAD)
This code supports the NeurIPS paper 'Implicit Deep Adaptive Design: Policy-Based Experimental Design without Likelihoods'.
@article{ivanova2021implicit,
title={Implicit Deep Adaptive Design: Policy-Based Experimental Design without Likelihoods},
author={Ivanova, Desi R. and Foster, Adam and Kleinegesse, Steven and Gutmann, Michael and Rainforth, Tom},
journal={Advances in Neural Information Processing Systems (NeurIPS)},
year={2021}
}
Computing infrastructure requirements
We have tested this codebase on Linux (Ubuntu x86_64) and MacOS (Big Sur v11.2.3) with Python 3.8. To train iDAD networks, we recommend the use of a GPU. We used one GeForce RTX 3090 GPU on a machine with 126 GiB of CPU memory and 40 CPU cores.
Installation
- Ensure that Python and
condaare installed. - Create and activate a new
condavirtual environment as follows
conda create -n idad_code
conda activate idad_code
- Install the correct version of PyTorch, following the instructions at pytorch.org. For our experiments we used
torch==1.8.0with CUDA version 11.1. - Install the remaining package requirements using
pip install -r requirements.txt. - Install the torchsde package from its repository:
pip install git+https://github.com/google-research/torchsde.git.
MLFlow
We use mlflow to log metric and store network parameters. Each experiment run is stored in a directory mlruns which will be created automatically. Each experiment is assigned a numerical
and each run gets a unique
. The iDAD networks will be saved in ./mlruns/
, which will be printed at the end of each training run.
Location Finding Experiment
To train an iDAD network with the InfoNCE bound to locate 2 sources in 2D, using the approach in the paper, execute the command
python3 location_finding.py \
--num-steps 100000 \
--num-experiments=10 \
--physical-dim 2 \
--num-sources 2 \
--lr 0.0005 \
--num-experiments 10 \
--encoding-dim 64 \
--hidden-dim 512 \
--mi-estimator InfoNCE \
--device <DEVICE>
To train an iDAD network with the NWJ bound, using the approach in the paper, execute the command
python3 location_finding.py \
--num-steps 100000 \
--num-experiments=10 \
--physical-dim 2 \
--num-sources 2 \
--lr 0.0005 \
--num-experiments 10 \
--encoding-dim 64 \
--hidden-dim 512 \
--mi-estimator NWJ \
--device <DEVICE>
To run the static MINEBED baseline, use the following
python3 location_finding.py \
--num-steps 100000 \
--physical-dim 2 \
--num-sources 2 \
--lr 0.0001 \
--num-experiments 10 \
--encoding-dim 8 \
--hidden-dim 512 \
--design-arch static \
--critic-arch cat \
--mi-estimator NWJ \
--device <DEVICE>
To run the static SG-BOED baseline, use the following
python3 location_finding.py \
--num-steps 100000 \
--physical-dim 2 \
--num-sources 2 \
--lr 0.0005 \
--num-experiments 10 \
--encoding-dim 8 \
--hidden-dim 512 \
--design-arch static \
--critic-arch cat \
--mi-estimator InfoNCE \
--device <DEVICE>
To run the adaptive (explicit likelihood) DAD baseline, use the following
python3 location_finding.py \
--num-steps 100000 \
--physical-dim 2 \
--num-sources 2 \
--lr 0.0005 \
--num-experiments 10 \
--encoding-dim 32 \
--hidden-dim 512 \
--mi-estimator sPCE \
--design-arch sum \
--device <DEVICE>
To evaluate the resulting networks eun the following command
python3 eval_sPCE.py --experiment-id <ID>
To evaluate a random design baseline (requires no pre-training):
python3 baselines_locfin_nontrainable.py \
--policy random \
--physical-dim 2 \
--num-experiments-to-perform 5 10 \
--device <DEVICE>
To run the variational baseline (note: it takes a very long time), run:
python3 baselines_locfin_variational.py \
--num-histories 128 \
--num-experiments 10 \
--physical-dim 2 \
--lr 0.001 \
--num-steps 5000\
--device <DEVICE>
Copy path_to_artifact and pass it to the evaluation script:
python3 eval_sPCE_from_source.py \
--path-to-artifact <path_to_artifact> \
--num-experiments-to-perform 5 10 \
--device <DEVICE>
Pharmacokinetic Experiment
To train an iDAD network with the InfoNCE bound, using the approach in the paper, execute the command
python3 pharmacokinetic.py \
--num-steps 100000 \
--lr 0.0001 \
--num-experiments 5 \
--encoding-dim 32 \
--hidden-dim 512 \
--mi-estimator InfoNCE \
--device <DEVICE>
To train an iDAD network with the NWJ bound, using the approach in the paper, execute the command
python3 pharmacokinetic.py \
--num-steps 100000 \
--lr 0.0001 \
--num-experiments 5 \
--encoding-dim 32 \
--hidden-dim 512 \
--mi-estimator NWJ \
--gamma 0.5 \
--device <DEVICE>
To run the static MINEBED baseline, use the following
python3 pharmacokinetic.py \
--num-steps 100000 \
--lr 0.001 \
--num-experiments 5 \
--encoding-dim 8 \
--hidden-dim 512 \
--design-arch static \
--critic-arch cat \
--mi-estimator NWJ \
--device <DEVICE>
To run the static SG-BOED baseline, use the following
python3 pharmacokinetic.py \
--num-steps 100000 \
--lr 0.0005 \
--num-experiments 5 \
--encoding-dim 8 \
--hidden-dim 512 \
--design-arch static \
--critic-arch cat \
--mi-estimator InfoNCE \
--device <DEVICE>
To run the adaptive (explicit likelihood) DAD baseline, use the following
python3 pharmacokinetic.py \
--num-steps 100000 \
--lr 0.0001 \
--num-experiments 5 \
--encoding-dim 32 \
--hidden-dim 512 \
--mi-estimator sPCE \
--design-arch sum \
--device <DEVICE>
To evaluate the resulting networks run the following command
python3 eval_sPCE.py --experiment-id <ID>
To evaluate a random design baseline (requires no pre-training):
python3 baselines_pharmaco_nontrainable.py \
--policy random \
--num-experiments-to-perform 5 10 \
--device <DEVICE>
To evaluate an equal interval baseline (requires no pre-training):
python3 baselines_pharmaco_nontrainable.py \
--policy equal_interval \
--num-experiments-to-perform 5 10 \
--device <DEVICE>
To run the variational baseline (note: it takes a very long time), run:
python3 baselines_pharmaco_variational.py \
--num-histories 128 \
--num-experiments 10 \
--lr 0.001 \
--num-steps 5000 \
--device <DEVICE>
Copy path_to_artifact and pass it to the evaluation script:
python3 eval_sPCE_from_source.py \
--path-to-artifact <path_to_artifact> \
--num-experiments-to-perform 5 10 \
--device <DEVICE>
SIR experiment
For the SIR experiments, please first generate an initial training set and a test set:
python3 epidemic_simulate_data.py \
--num-samples=100000 \
--device <DEVICE>
To train an iDAD network with the InfoNCE bound, using the approach in the paper, execute the command
python3 epidemic.py \
--num-steps 100000 \
--num-experiments 5 \
--lr 0.0005 \
--hidden-dim 512 \
--encoding-dim 32 \
--mi-estimator InfoNCE \
--design-transform ts \
--device <DEVICE>
To train an iDAD network with the NWJ bound, execute the command
python3 epidemic.py \
--num-steps 100000 \
--num-experiments 5 \
--lr 0.0005 \
--hidden-dim 512 \
--encoding-dim 32 \
--mi-estimator NWJ \
--design-transform ts \
--device <DEVICE>
To run the static SG-BOED baseline, run
python3 epidemic.py \
--num-steps 100000 \
--num-experiments 5 \
--lr 0.005 \
--hidden-dim 512 \
--encoding-dim 32 \
--design-arch static \
--critic-arch cat \
--design-transform iid \
--mi-estimator InfoNCE \
--device <DEVICE>
To run the static MINEBED baseline, run
python3 epidemic.py \
--num-steps 100000 \
--num-experiments 5 \
--lr 0.001 \
--hidden-dim 512 \
--encoding-dim 32 \
--design-arch static \
--critic-arch cat \
--design-transform iid \
--mi-estimator NWJ \
--device <DEVICE>
To train a critic with random designs (to evaluate the random design baseline):
python3 epidemic.py \
--num-steps 100000 \
--num-experiments 5 \
--lr 0.005 \
--hidden-dim 512 \
--encoding-dim 32 \
--design-arch random \
--critic-arch cat \
--design-transform iid \
--device <DEVICE>
To train a critic with equal interval designs, which is then used to evaluate the equal interval baseline, run the following
python3 epidemic.py \
--num-steps 100000 \
--num-experiments 5 \
--lr 0.001 \
--hidden-dim 512 \
--encoding-dim 32 \
--design-arch equal_interval \
--critic-arch cat \
--design-transform iid \
--device <DEVICE>
Finally, to evaluate the different methods, run
python3 eval_epidemic.py \
--experiment-id <ID> \
--device <DEVICE>
80 Nov 20, 2022
0 Oct 18, 2021
66 Jul 11, 2022
106 Jan 04, 2023
0 Dec 01, 2021
2 Nov 08, 2022
16 Nov 14, 2022
98 Dec 07, 2022
2.4k Jan 04, 2023
17 Aug 04, 2022
386 Dec 05, 2022
7 Oct 02, 2022
273 Dec 17, 2022
19 Dec 13, 2022
2 Sep 11, 2022
305 Dec 29, 2022
6.7k Dec 31, 2022
23 Dec 13, 2022
5.1k Dec 30, 2022
48 Dec 26, 2022