Active Learning Experiment

In this module, we find all the utilities to do active learning. Baal takes care of the dataset split between labelled and unlabelled examples. It also takes care of the active learning loop:

1. Train the model on the label set.
2. Evaluate the model on the evaluation set.
3. Predict on the unlabelled examples.
4. Label the most uncertain examples.
5. Stop the experiment if finished.

Example

from baal.active.dataset import ActiveLearningDataset
from baal.experiments.base import ActiveLearningExperiment
al_dataset = ActiveLearningDataset(your_dataset)

# To start, we can select 1000 random examples to be labelled
al_dataset.label_randomly(1000)

experiment = ActiveLearningExperiment(
    trainer=..., # Huggingface or ModelWrapper to train
    al_dataset=al_dataset, # Active learning dataset
    eval_dataset=..., # Evaluation Dataset
    heuristic=BALD(), # Uncertainty heuristic to use
    query_size=100, # How many items to label per round.
    iterations=20, # How many MC sampling to perform per item.
    pool_size=None, # Optionally limit the size of the unlabelled pool.
    criterion=None # Stopping criterion for the experiment.
)
experiment.start()

API

baal.experiments.base.ActiveLearningExperiment

Experiment manager for Baal.

Takes care of

1. Train the model on the label set.
2. Evaluate the model on the evaluation set.
3. Predict on the unlabelled examples.
4. Label the most uncertain examples.
5. Stop the experiment if finished.

Parameters:

Name	Type	Description	Default
trainer	Union[ModelWrapper, BaalTransformersTrainer]	Huggingface or ModelWrapper to train	required
al_dataset	ActiveLearningDataset	Active learning dataset	required
eval_dataset	Dataset	Evaluation Dataset	required
heuristic	AbstractHeuristic	Uncertainty heuristic to use	required
query_size	int	How many items to label per round.	100
iterations	int	How many MC sampling to perform per item.	20
pool_size	Optional[int]	Optionally limit the size of the unlabelled pool.	None
criterion	Optional[StoppingCriterion]	Stopping criterion for the experiment.	None

Source code in baal/experiments/base.py

class ActiveLearningExperiment:
    """Experiment manager for Baal.

    Takes care of:
        1. Train the model on the label set.
        2. Evaluate the model on the evaluation set.
        3. Predict on the unlabelled examples.
        4. Label the most uncertain examples.
        5. Stop the experiment if finished.

    Args:
        trainer: Huggingface or ModelWrapper to train
        al_dataset: Active learning dataset
        eval_dataset: Evaluation Dataset
        heuristic: Uncertainty heuristic to use
        query_size: How many items to label per round.
        iterations: How many MC sampling to perform per item.
        pool_size: Optionally limit the size of the unlabelled pool.
        criterion: Stopping criterion for the experiment.
    """

    def __init__(
        self,
        trainer: Union[ModelWrapper, "BaalTransformersTrainer"],
        al_dataset: ActiveLearningDataset,
        eval_dataset: Dataset,
        heuristic: AbstractHeuristic,
        query_size: int = 100,
        iterations: int = 20,
        pool_size: Optional[int] = None,
        criterion: Optional[StoppingCriterion] = None,
    ):
        self.al_dataset = al_dataset
        self.eval_dataset = eval_dataset
        self.heuristic = heuristic
        self.query_size = query_size
        self.iterations = iterations
        self.criterion = criterion or LabellingBudgetStoppingCriterion(
            al_dataset, labelling_budget=al_dataset.n_unlabelled
        )
        self.pool_size = pool_size
        self.adapter = self._get_adapter(trainer)

    def start(self):
        records = []
        _start = len(self.al_dataset)
        if _start == 0:
            raise ValueError(
                "No item labelled in the training set."
                " Did you run `ActiveLearningDataset.label_randomly`?"
            )
        for _ in tqdm(
            itertools.count(start=0),  # Infinite counter to rely on Criterion
            desc="Active Experiment",
            # Upper bound estimation.
            total=np.round(self.al_dataset.n_unlabelled // self.query_size),
        ):
            self.adapter.reset_weights()
            train_metrics = self.adapter.train(self.al_dataset)
            eval_metrics = self.adapter.evaluate(
                self.eval_dataset, average_predictions=self.iterations
            )
            pool = self._get_pool()
            ranks, uncertainty = self.heuristic.get_ranks(
                self.adapter.predict(pool, iterations=self.iterations)
            )
            self.al_dataset.label(ranks[: self.query_size])
            records.append({**train_metrics, **eval_metrics})
            if self.criterion.should_stop(eval_metrics, uncertainty):
                log.info("Experiment complete", num_labelled=len(self.al_dataset) - _start)
                return records

    def _get_adapter(
        self, trainer: Union[ModelWrapper, "BaalTransformersTrainer"]
    ) -> FrameworkAdapter:
        if isinstance(trainer, ModelWrapper):
            return ModelWrapperAdapter(trainer)
        elif TRANSFORMERS_AVAILABLE and isinstance(trainer, BaalTransformersTrainer):
            return TransformersAdapter(trainer)
        raise ValueError(
            f"{type(trainer)} is not a supported trainer."
            " Baal supports ModelWrapper and BaalTransformersTrainer"
        )

    def _get_pool(self):
        if self.pool_size is None:
            return self.al_dataset.pool
        pool = self.al_dataset.pool
        if len(pool) < self.pool_size:
            return pool
        indices = np.random.choice(len(pool), min(len(pool), self.pool_size), replace=False)
        return Subset(pool, indices)