Skip to main content

An implementation of the BatchBALD algorithm

Project description

BatchBALD Redux

Clean reimplementation of "BatchBALD: Efficient and Diverse Batch Acquisition for Deep Bayesian Active Learning"

For an introduction & more information, see The paper can be found at

The original implementation used in the paper is available at

We are grateful for fastai's nbdev which is powering this package.


pip install batchbald_redux


BatchBALD is an algorithm and acquisition function for Active Learning in a Bayesian setting using BNNs and MC dropout.

The aquisition function is the mutual information between the joint of a candidate batch and the model parameters $\omega$:

{% raw %} $$a_{\text{BatchBALD}}((y_b)_B) = I[(y_b)_B;\omega]$$ {% endraw %}

The best candidate batch is one that maximizes this acquisition function.

In the paper, we show that this function satisfies sub-modularity, which provides us an optimality guarantee for a greedy algorithm. The candidate batch is selected using greedy expansion.

Joint entropies are hard to estimate and, for everything to work, one also has to use consistent MC dropout, which keeps a set of dropout masks fixed while scoring the pool set.

To aid reproducibility and baseline reproduction, we provide this simpler and clearer reimplementation.

How to use

We provide a simple example experiment that uses this package here.

To get a candidate batch using BatchBALD, we provide a simple API in batchbald_redux.batchbald:

from nbdev.showdoc import *
from batchbald_redux.batchbald import get_batchbald_batch



get_batchbald_batch(logits_N_K_C:Tensor, batch_size:int, num_samples:int, dtype=None, device=None)

We also provide a simple implementation of consistent MC dropout in batchbald_redux.consistent_mc_dropout.

Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Files for batchbald-redux, version 1.0.1
Filename, size File type Python version Upload date Hashes
Filename, size batchbald_redux-1.0.1-py3-none-any.whl (16.7 kB) File type Wheel Python version py3 Upload date Hashes View
Filename, size batchbald_redux-1.0.1.tar.gz (17.0 kB) File type Source Python version None Upload date Hashes View

Supported by

Pingdom Pingdom Monitoring Google Google Object Storage and Download Analytics Sentry Sentry Error logging AWS AWS Cloud computing DataDog DataDog Monitoring Fastly Fastly CDN DigiCert DigiCert EV certificate StatusPage StatusPage Status page