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>
3 Oct 09, 2022
9 Dec 23, 2022
28 Jan 08, 2023
38 Dec 26, 2022
5 Dec 22, 2022
76 Nov 23, 2022
922 Dec 21, 2022
160 Dec 19, 2022
1 Jan 05, 2022
34 Sep 10, 2022
125 Dec 24, 2022
53 Dec 26, 2022
24 Mar 23, 2022
6 Oct 08, 2022
1 Feb 09, 2022
17 Oct 17, 2022
64 Dec 24, 2022
1 Nov 13, 2021
1 Dec 10, 2021
53 Dec 17, 2022