Scalable Batch Acquisition for Deep Bayesian Active Learning

This is a PyTorch implementation of the SDM 2023 paper Scalable Batch Acquisition for Deep Bayesian Active Learning. Our work present a new Bayesian active learning algorithm called Large BatchBALD aka LBB (and its stochastic extension Power Large BatchBALD aka PLBB), which gives a well-grounded approximation to the BatchBALD method and aims to achieve comparable quality while being more computationally efficient.

Install

pip install batchbald_redux

Train-evaluation examples

FMNIST with MC-dropout:

python sampling_train.py --dataset_name='FMNIST' --model_name='CNN_MC_RMNIST' --uns_type='MC' --algs PLBB PBALD Rand LBB BALD BB MaxProb --random_seeds 42 227 346 684 920 --acq_batch_size=10 --num_init_samples=20 --max_train_samples=500

RCIFAR-10 with deep ensembles:

python epochs_train.py --dataset='RCIFAR10' --model_name='ResNet-18' --optimizer_name='SGD' --uns_type='ENS' --algs PLBB PBALD Rand LBB BALD MaxProb --random_seeds 42 227 346 684 920 --acq_batch_size=100 --train_batch_size=100 --num_init_samples=2000 --max_train_samples=10000 --num_epochs=50

Experimental setup

All code related to the BALD and BatchBALD algorithms and the corresponding modules are taken from the batchbald_redux repository. All datasets and implemented algorithms are placed in the batchbald_redux/ directory. Train-evaluation files are placed in the main directory, results are saved in the results/ directory with the corresponding config. Several training options are available: sample-wise training via sampling_train.py and regular epoch-wise training via epochs_train.py. There are also several options for uncertainty estimation: MC-dropout and deep ensembles (available through the training arguments).

Available active learning algorithms:

• BALD
• PowerBALD
• BatchBALD
• Large BatchBALD (Ours)
• Power Large BatchBALD (Ours)
• MaxProb
• Entropy sampling
• Random sampling

Datasets:

• MNIST-based: MNIST, RMNIST, FMNIST, EMNIST, KMNIST
• CIFAR-based: CIFAR-10, CIFAR-100, RCIFAR-10, RCIFAR-100
• Others: SVHN, AG News (text)

Citation

@article{rubashevskii2023sbadbal,
  title={Scalable Batch Acquisition for Deep Bayesian Active Learning},
  author={Rubashevskii, Aleksandr and Kotova, Daria and Panov, Maxim},
  journal={arXiv preprint arXiv:2301.05490},
  year={2023}
}

Big thanks to batchbald_redux, our code is partially borrowed from them.