Repository for bayesian_active_learning, an implementation of "Deep Bayesian Active Learning with Image Data" using Haiku and Optax.
The repository has a module bayesian_active_learning which contains:
acquisition_functions.py: acquistion functions and sampling methodology for active learning proceduredata_utils.py: Simple classes for loading datasets from numpy arraysexperiment.py: Contains end to end procedure for active learning setuplosses.py: loss functions (only classification loss implemented for now)metrics.py: simple methods to evaluate jax models given parameters and a dataloadermodels.py: Bayesian convnet implementation and Haiku transformation methodtraining.py:fitmethod for training haiku models with optaxutils.py: miscellaneous utility functions
A series of notebooks in the notebooks folder investigate different aspects of the paper:
00_weight_decay_tuning.ipynb: notebook to finetune weight decay hyperparameter on validation dataset01_standard_experiment.ipynb: Partial reconstruction of Figure 1 from the paper, using a subselection of the acquisition functions used in the paper with a relatively large batch size
02_reduced_batch_size.ipynb: An experiment using a much lower acquisition size per iteration (20 rather than 100), since using the standard BALD acquisition function can have degraded performance under large acquisition sizes (see "BatchBALD: Efficient and Diverse Batch Acquisition for Deep Bayesian Active Learning" )
03_modified_pool_set.ipynb: An experiment where the pool set (unlabelled images that are selected during active learning) is made artifically unbalanced, by removing 75% of the "2"s in the set.
04_further_extensions.ipynb: a short notebook detailing other ideas for future work.