API Reference

ModelWrapper

class baal.ModelWrapper(model, criterion, replicate_in_memory=True)[source]

Wrapper created to ease the training/testing/loading.

Parameters
  • model (nn.Module) – The model to optimize.

  • criterion (Callable) – A loss function.

  • replicate_in_memory (bool) – Replicate in memory optional.

Methods

add_metric(name, initializer)

Add a baal.utils.metric.Metric to the Model.

eval()

Set the model in eval mode.

get_params()

Return the parameters to optimize.

load_state_dict(state_dict[, strict])

Load the model with state_dict.

predict_on_batch(data[, iterations, cuda])

Get the model’s prediction on a batch.

predict_on_dataset(dataset, batch_size, …)

Use the model to predict on a dataset iterations time.

predict_on_dataset_generator(dataset, …[, …])

Use the model to predict on a dataset iterations time.

reset_all()

Reset all resetable layers.

reset_fcs()

Reset all torch.nn.Linear layers.

state_dict()

Get the state dict(s).

test_on_batch(data, target[, cuda, …])

Test the current model on a batch.

test_on_dataset(dataset, batch_size, use_cuda)

Test the model on a Dataset dataset.

train()

Set the model in train mode.

train_and_test_on_datasets(train_dataset, …)

Train and test the model on both Dataset train_dataset, test_dataset.

train_on_batch(data, target, optimizer[, …])

Train the current model on a batch using optimizer.

train_on_dataset(dataset, optimizer, …[, …])

Train for epoch epochs on a Dataset `dataset.

add_metric(name: str, initializer: Callable)[source]

Add a baal.utils.metric.Metric to the Model.

Parameters
  • name (str) – name of the metric.

  • initializer (Callable) – lambda to initialize a new instance of a baal.utils.metrics.Metric object.

eval()[source]

Set the model in eval mode.

get_params()[source]

Return the parameters to optimize.

Returns

Config for parameters.

load_state_dict(state_dict, strict=True)[source]

Load the model with state_dict.

predict_on_batch(data, iterations=1, cuda=False)[source]

Get the model’s prediction on a batch.

Parameters
  • data (Tensor) – The model input.

  • iterations (int) – Number of prediction to perform.

  • cuda (bool) – Use CUDA or not.

Returns

Tensor, the loss computed from the criterion.

shape = {batch_size, nclass, n_iteration}.

Raises

Raises RuntimeError if CUDA rans out of memory during data replication.

predict_on_dataset(dataset: torch.utils.data.Dataset, batch_size: int, iterations: int, use_cuda: bool, workers: int = 4, collate_fn: Optional[Callable] = None, half=False, verbose=True)[source]

Use the model to predict on a dataset iterations time.

Parameters
  • dataset (Dataset) – Dataset to predict on.

  • batch_size (int) – Batch size to use during prediction.

  • iterations (int) – Number of iterations per sample.

  • use_cuda (bool) – Use CUDA or not.

  • workers (int) – Number of workers to use.

  • collate_fn (Optional[Callable]) – The collate function to use.

  • half (bool) – If True use half precision.

  • verbose (bool) – If True use tqdm to show progress.

Notes

The “batch” is made of batch_size * iterations samples.

Returns

Array [n_samples, n_outputs, …, n_iterations].

predict_on_dataset_generator(dataset: torch.utils.data.Dataset, batch_size: int, iterations: int, use_cuda: bool, workers: int = 4, collate_fn: Optional[Callable] = None, half=False, verbose=True)[source]

Use the model to predict on a dataset iterations time.

Parameters
  • dataset (Dataset) – Dataset to predict on.

  • batch_size (int) – Batch size to use during prediction.

  • iterations (int) – Number of iterations per sample.

  • use_cuda (bool) – Use CUDA or not.

  • workers (int) – Number of workers to use.

  • collate_fn (Optional[Callable]) – The collate function to use.

  • half (bool) – If True use half precision.

  • verbose (bool) – If True use tqdm to display progress

Notes

The “batch” is made of batch_size * iterations samples.

Returns

Generators [batch_size, n_classes, …, n_iterations].

reset_all()[source]

Reset all resetable layers.

reset_fcs()[source]

Reset all torch.nn.Linear layers.

state_dict()[source]

Get the state dict(s).

test_on_batch(data: torch.Tensor, target: torch.Tensor, cuda: bool = False, average_predictions: int = 1)[source]

Test the current model on a batch.

Parameters
  • data (Tensor) – The model input.

  • target (Tensor) – The ground truth.

  • cuda (bool) – Use CUDA or not.

  • average_predictions (int) – The number of predictions to average to compute the test loss.

Returns

Tensor, the loss computed from the criterion.

test_on_dataset(dataset: torch.utils.data.Dataset, batch_size: int, use_cuda: bool, workers: int = 4, collate_fn: Optional[Callable] = None, average_predictions: int = 1)[source]

Test the model on a Dataset dataset.

Parameters
  • dataset (Dataset) – Dataset to evaluate on.

  • batch_size (int) – Batch size used for evaluation.

  • use_cuda (bool) – Use Cuda or not.

  • workers (int) – Number of workers to use.

  • collate_fn (Optional[Callable]) – The collate function to use.

  • average_predictions (int) – The number of predictions to average to compute the test loss.

Returns

Average loss value over the dataset.

train()[source]

Set the model in train mode.

train_and_test_on_datasets(train_dataset: torch.utils.data.Dataset, test_dataset: torch.utils.data.Dataset, optimizer: torch.optim.Optimizer, batch_size: int, epoch: int, use_cuda: bool, workers: int = 4, collate_fn: Optional[Callable] = None, regularizer: Optional[Callable] = None, return_best_weights=False, patience=None, min_epoch_for_es=0)[source]

Train and test the model on both Dataset train_dataset, test_dataset.

Parameters
  • train_dataset (Dataset) – Dataset to train on.

  • test_dataset (Dataset) – Dataset to evaluate on.

  • optimizer (Optimizer) – Optimizer to use during training.

  • batch_size (int) – Batch size used.

  • epoch (int) – Number of epoch to train on.

  • use_cuda (bool) – Use Cuda or not.

  • workers (int) – Number of workers to use.

  • collate_fn (Optional[Callable]) – The collate function to use.

  • regularizer (Optional[Callable]) – The loss regularization for training.

  • return_best_weights (bool) – If True, will keep the best weights and return them.

  • patience (Optional[int]) – If provided, will use early stopping to stop after patience epoch without improvement.

  • min_epoch_for_es (int) – Epoch at which the early stopping starts.

Returns

History and best weights if required.

train_on_batch(data, target, optimizer, cuda=False, regularizer: Optional[Callable] = None)[source]

Train the current model on a batch using optimizer.

Parameters
  • data (Tensor) – The model input.

  • target (Tensor) – The ground truth.

  • optimizer (optim.Optimizer) – An optimizer.

  • cuda (bool) – Use CUDA or not.

  • regularizer (Optional[Callable]) – The loss regularization for training.

Returns

Tensor, the loss computed from the criterion.

train_on_dataset(dataset, optimizer, batch_size, epoch, use_cuda, workers=4, collate_fn: Optional[Callable] = None, regularizer: Optional[Callable] = None)[source]

Train for epoch epochs on a Dataset `dataset.

Parameters
  • dataset (Dataset) – Pytorch Dataset to be trained on.

  • optimizer (optim.Optimizer) – Optimizer to use.

  • batch_size (int) – The batch size used in the DataLoader.

  • epoch (int) – Number of epoch to train for.

  • use_cuda (bool) – Use cuda or not.

  • workers (int) – Number of workers for the multiprocessing.

  • collate_fn (Optional[Callable]) – The collate function to use.

  • regularizer (Optional[Callable]) – The loss regularization for training.

Returns

The training history.

Active learning functionality

class baal.active.ActiveLearningDataset(*args: Any, **kwargs: Any)[source]

A dataset that allows for active learning.

Parameters
  • dataset (torch.data.Dataset) – The baseline dataset.

  • labelled (Union[np.ndarray, torch.Tensor]) – An array/tensor that acts as a boolean mask which is True for every data point that is labelled, and False for every data point that is not labelled.

  • make_unlabelled (Callable) – The function that returns an unlabelled version of a datum so that it can still be used in the DataLoader.

  • random_state (None, int, RandomState) – Set the random seed for label_randomly().

  • pool_specifics (Optional[Dict]) – Attributes to set when creating the pool. Useful to remove data augmentation.

Attributes
n_labelled

The number of labelled data points.

n_unlabelled

The number of unlabelled data points.

pool

Returns a new Dataset made from unlabelled samples.

Methods

ActiveIter(aldataset)

Iterator over an ActiveLearningDataset.

__call__(*args, **kwargs)

Call self as a function.

check_dataset_can_label()

Check if a dataset can be labelled.

get_raw(idx)

Get a datapoint from the underlying dataset.

is_labelled(idx)

Check if a datapoint is labelled.

label(index[, value])

Label data points.

label_randomly([n])

Label n data-points randomly.

load_state_dict(state_dict)

Load the labelled map and random_state with give state_dict.

reset_labeled()

Reset the label map.

state_dict()

Return the state_dict, ie.

class ActiveIter(aldataset)[source]

Iterator over an ActiveLearningDataset.

check_dataset_can_label()[source]

Check if a dataset can be labelled.

Returns

Whether the dataset’s label can be modified or not.

Notes

To be labelled, a dataset needs a method label with definition: label(self, idx, value) where value is the label for indice idx.

get_raw(idx: int) None[source]

Get a datapoint from the underlying dataset.

is_labelled(idx: int) bool[source]

Check if a datapoint is labelled.

label(index: Union[list, int], value: Optional[Any] = None) None[source]

Label data points. The index should be relative to the pool, not the overall data.

Parameters
  • index (Union[list,int]) – one or many indices to label.

  • value (Optional[Any]) – The label value. If not provided, no modification to the underlying dataset is done.

label_randomly(n: int = 1) None[source]

Label n data-points randomly.

Parameters

n (int) – Number of samples to label.

load_state_dict(state_dict)[source]

Load the labelled map and random_state with give state_dict.

property n_labelled

The number of labelled data points.

property n_unlabelled

The number of unlabelled data points.

property pool: torch.utils.data.Dataset

Returns a new Dataset made from unlabelled samples.

Raises

ValueError if a pool specific attribute cannot be set.

reset_labeled()[source]

Reset the label map.

state_dict()[source]

Return the state_dict, ie. the labelled map and random_state.

class baal.active.ActiveLearningLoop(dataset: baal.active.dataset.ActiveLearningDataset, get_probabilities: Callable, heuristic: baal.active.heuristics.heuristics.AbstractHeuristic = <baal.active.heuristics.heuristics.Random object>, ndata_to_label: int = 1, max_sample=-1, **kwargs)[source]

Object that perform the active learning iteration.

Parameters
  • dataset (ActiveLearningDataset) – Dataset with some sample already labelled.

  • get_probabilities (Function) – Dataset -> **kwargs -> ndarray [n_samples, n_outputs, n_iterations].

  • heuristic (Heuristic) – Heuristic from baal.active.heuristics.

  • ndata_to_label (int) – Number of sample to label per step.

  • max_sample (int) – Limit the number of sample used (-1 is no limit).

  • **kwargs – Parameters forwarded to get_probabilities.

Methods

step([pool])

Perform an active learning step.

step(pool=None) bool[source]

Perform an active learning step.

Parameters

pool (iterable) – dataset pool indices.

Returns

boolean, Flag indicating if we continue training.

class baal.active.FileDataset(*args: Any, **kwargs: Any)[source]

Dataset object that load the files and apply a transformation.

Parameters
  • files (List[str]) – The files.

  • lbls (List[Any]) – The labels, -1 indicates that the label is unknown.

  • transform (Optional[Callable]) – torchvision.transform pipeline.

  • target_transform (Optional[Callable]) – Function that modifies the target.

  • image_load_fn (Optional[Callable]) – Function that loads the image, by default uses PIL.

  • seed (Optional[int]) – Will set a seed before and between DA.

Methods

__call__(*args, **kwargs)

Call self as a function.

label(idx, lbl)

Label the sample idx with lbl.

get_kwargs

label(idx: int, lbl: Any)[source]

Label the sample idx with lbl.

Parameters
  • idx (int) – The sample index.

  • lbl (Any) – The label to assign.

Calibration Wrapper

class baal.calibration.DirichletCalibrator(wrapper: baal.modelwrapper.ModelWrapper, num_classes: int, lr: float, reg_factor: float, mu: Optional[float] = None)[source]

Adding a linear layer to a classifier model after the model is trained and train this new layer until convergence. Together with the linear layer, the model is now calibrated. Source: https://arxiv.org/abs/1910.12656 Code inspired from: https://github.com/dirichletcal/experiments_neurips

References

@article{kullbeyond,
title={Beyond temperature scaling: Obtaining well-calibrated multi-class

probabilities with Dirichlet calibration Supplementary material},

author={Kull, Meelis and Perello-Nieto,
Miquel and K{“a}ngsepp, Markus and Silva Filho,

Telmo and Song, Hao and Flach, Peter}

}

Parameters
  • wrapper (ModelWrapper) – Provides training and testing methods.

  • num_classes (int) – Number of classes in classification task.

  • lr (float) – Learning rate.

  • reg_factor (float) – Regularization factor for the linear layer weights.

  • mu (float) – Regularization factor for the linear layer biases. If not given, will be initialized by “l”.

Attributes
calibrated_model
metrics

Methods

calibrate(train_set, test_set, batch_size, …)

Training the linear layer given a training set and a validation set.

l2_reg()

Using trainable layer’s parameters for l2 regularization.

calibrate(train_set: torch.utils.data.Dataset, test_set: torch.utils.data.Dataset, batch_size: int, epoch: int, use_cuda: bool, double_fit: bool = False, **kwargs)[source]

Training the linear layer given a training set and a validation set. The training set should be different from what model is trained on.

Parameters
  • train_set (Dataset) – The training set.

  • test_set (Dataset) – The validation set.

  • batch_size (int) – Batch size used.

  • epoch (int) – Number of epochs to train the linear layer for.

  • use_cuda (bool) – If “True”, will use GPU.

  • double_fit (bool) – If “True” would fit twice on the train set.

  • kwargs (dict) – Rest of parameters for baal.ModelWrapper.train_and_test_on_dataset().

Returns

List of loss values for each epoch. model.state_dict (dict): Model weights.

Return type

loss_history (list[float])

l2_reg()[source]

Using trainable layer’s parameters for l2 regularization.

Returns

The regularization term for the linear layer.

Heuristics

class baal.active.heuristics.AbstractHeuristic(shuffle_prop=0.0, reverse=False, reduction='none')[source]

Abstract class that defines a Heuristic.

Parameters
  • shuffle_prop (float) – shuffle proportion.

  • reverse (bool) – True if the most uncertain sample has the highest value.

  • reduction (Union[str, Callable]) – Reduction used after computing the score.

Methods

__call__(predictions)

Rank the predictions according to their uncertainties.

compute_score(predictions)

Compute the score according to the heuristic.

get_ranks(predictions)

Rank the predictions according to their uncertainties.

get_uncertainties(predictions)

Get the uncertainties.

get_uncertainties_generator(predictions)

Compute the score according to the heuristic.

reorder_indices(scores)

Order indices given their uncertainty score.

compute_score(predictions)[source]

Compute the score according to the heuristic.

Parameters

predictions (ndarray) – Array of predictions

Returns

Array of scores.

get_ranks(predictions)[source]

Rank the predictions according to their uncertainties.

Parameters

predictions (ndarray) – [batch_size, C, …, Iterations]

Returns

Ranked index according to the uncertainty (highest to lowes).

get_uncertainties(predictions)[source]

Get the uncertainties.

Parameters

predictions (ndarray) – Array of predictions

Returns

Array of uncertainties

get_uncertainties_generator(predictions)[source]

Compute the score according to the heuristic.

Parameters

predictions (Iterable) – Generator of predictions

Raises

ValueError if the generator is empty.

Returns

Array of scores.

reorder_indices(scores)[source]

Order indices given their uncertainty score.

Parameters

scores (ndarray/ List[ndarray]) – Array of uncertainties or list of arrays.

Returns

ordered index according to the uncertainty (highest to lowes).

Raises

ValueError if scores is not uni-dimensional.

class baal.active.heuristics.BALD(shuffle_prop=0.0, reduction='none')[source]

Sort by the highest acquisition function value.

Parameters
  • shuffle_prop (float) – Amount of noise to put in the ranking. Helps with selection bias (default: 0.0).

  • reduction (Union[str, callable]) – function that aggregates the results (default: ‘none`).

References

https://arxiv.org/abs/1703.02910

Methods

__call__(predictions)

Rank the predictions according to their uncertainties.

compute_score(predictions)

Compute the score according to the heuristic.

get_ranks(predictions)

Rank the predictions according to their uncertainties.

get_uncertainties(predictions)

Get the uncertainties.

get_uncertainties_generator(predictions)

Compute the score according to the heuristic.

reorder_indices(scores)

Order indices given their uncertainty score.

class baal.active.heuristics.Random(shuffle_prop=0.0, reduction='none')[source]

Random heuristic.

Parameters
  • shuffle_prop (float) – UNUSED

  • reduction (Union[str, callable]) – UNUSED.

Methods

__call__(predictions)

Rank the predictions according to their uncertainties.

compute_score(predictions)

Compute the score according to the heuristic.

get_ranks(predictions)

Rank the predictions according to their uncertainties.

get_uncertainties(predictions)

Get the uncertainties.

get_uncertainties_generator(predictions)

Compute the score according to the heuristic.

reorder_indices(predictions)

Order indices randomly.

class baal.active.heuristics.Entropy(shuffle_prop=0.0, reduction='none')[source]

Sort by the highest entropy.

Parameters
  • shuffle_prop (float) – Amount of noise to put in the ranking. Helps with selection bias (default: 0.0).

  • reduction (Union[str, callable]) – function that aggregates the results (default: none).

Methods

__call__(predictions)

Rank the predictions according to their uncertainties.

compute_score(predictions)

Compute the score according to the heuristic.

get_ranks(predictions)

Rank the predictions according to their uncertainties.

get_uncertainties(predictions)

Get the uncertainties.

get_uncertainties_generator(predictions)

Compute the score according to the heuristic.

reorder_indices(scores)

Order indices given their uncertainty score.

Pytorch Lightning Compatibility

class baal.utils.pytorch_lightning.ResetCallback(*args: Any, **kwargs: Any)[source]

Callback to reset the weights between active learning steps.

Parameters

weights (dict) – State dict of the model.

Notes

The weight should be deep copied beforehand.

Methods

on_train_start(trainer, module)

Will reset the module to its initial weights.

on_train_start(trainer, module)[source]

Will reset the module to its initial weights.

class baal.utils.pytorch_lightning.BaalTrainer(*args: Any, **kwargs: Any)[source]

Object that perform the training and active learning iteration.

Parameters
  • dataset (ActiveLearningDataset) – Dataset with some sample already labelled.

  • heuristic (Heuristic) – Heuristic from baal.active.heuristics.

  • ndata_to_label (int) – Number of sample to label per step.

  • max_sample (int) – Limit the number of sample used (-1 is no limit).

  • **kwargs – Parameters forwarded to get_probabilities and to pytorch_ligthning Trainer.__init__

Methods

predict_on_dataset([model, dataloader])

For documentation, see predict_on_dataset_generator

predict_on_dataset_generator([model, dataloader])

Predict on the pool loader.

predict_on_dataset(model=None, dataloader=None, *args, **kwargs)[source]

For documentation, see predict_on_dataset_generator

predict_on_dataset_generator(model=None, dataloader=None, *args, **kwargs)[source]

Predict on the pool loader.

Parameters
  • model – Model to be used in prediction. If None, will get the Trainer’s model.

  • dataloader (Optional[DataLoader]) – If provided, will predict on this dataloader. Otherwise, uses model.pool_dataloader().

Returns

Numpy arrays with all the predictions.

class baal.utils.pytorch_lightning.BaaLDataModule(*args: Any, **kwargs: Any)[source]

Methods

pool_dataloader()

Create Dataloader for the pool of unlabelled examples.

pool_dataloader()[source]

Create Dataloader for the pool of unlabelled examples.