Experiment API v2 (#296)

* WIP implementation

* Finish Adapter work, start test and editing experiments

* Add tests to HF

* Lot of documentation

* Mypy

* Lint

* Add test for Experiment.get_pool

* Update to autotqdm, improve example

* Improve base example

* Fix documentation

* Bump to 2.0.0

README.md

@@ -91,7 +91,7 @@ The framework consists of four main parts, as demonstrated in the flowchart belo
 - ActiveLearningLoop
 
 <p align="center">
- <img src="docs/research/literature/images/Baalscheme.svg">
+ <img src="docs/learn/literature/images/Baalscheme.svg">
 </p>
 
 To get started, wrap your dataset in our _[**ActiveLearningDataset**](baal/active/dataset.py)_ class. This will ensure
@@ -114,19 +114,19 @@ In conclusion, your script should be similar to this:
 dataset = ActiveLearningDataset(your_dataset)
 dataset.label_randomly(INITIAL_POOL) # label some data
 model = MCDropoutModule(your_model)
-model = ModelWrapper(model, args=TrainingArgs(...))
-active_loop = ActiveLearningLoop(dataset,
-  get_probabilities=model.predict_on_dataset,
-  heuristic=heuristics.BALD(),
-  iterations=20, # Number of MC sampling.
-  query_size=QUERY_SIZE) # Number of item to label.
-for al_step in range(N_ALSTEP):
- model.train_on_dataset(dataset)
- metrics = model.test_on_dataset(test_dataset)
- # Label the next most uncertain items.
- if not active_loop.step():
- # We're done!
- break
+wrapper = ModelWrapper(model, args=TrainingArgs(...))
+experiment = ActiveLearningExperiment(
+ trainer=wrapper, # Huggingface or ModelWrapper to train
+ al_dataset=dataset, # Active learning dataset
+ eval_dataset=test_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.
+)
+# The experiment will run until all items are labelled.
+metrics = experiment.start()
 ```
 
 For a complete experiment, see _[experiments/vgg_mcdropout_cifar10.py](experiments/vgg_mcdropout_cifar10.py)_ .

baal/active/active_loop.py

@@ -11,7 +11,7 @@
 from . import heuristics
 from .dataset import ActiveLearningDataset
 
-log = structlog.get_logger(__name__)
+log = structlog.get_logger("baal")
 pjoin = os.path.join
 
 

baal/active/dataset/pytorch_dataset.py

@@ -188,12 +188,12 @@ def label(self, index: Union[list, int], value: Optional[Any] = None) -> None:
  elif self.can_label and val is None:
  raise ValueError(
  """The dataset is able to label data, but no label was provided.
- If this is a research setting, please set the
+ If this is a learn setting, please set the
  `ActiveLearningDataset.can_label` to `False`.
  """
  )
  else:
- # Regular research usecase.
+ # Regular learn usecase.
  self.labelled_map[idx] = active_step
  if val is not None:
  warnings.warn(

baal/active/heuristics/heuristics_gpu.py

@@ -105,7 +105,11 @@ def predict_on_dataset(
  half=False,
  verbose=True,
  ):
- return super().predict_on_dataset(dataset, iterations, half, verbose).reshape([-1])
+ return (
+ super()
+ .predict_on_dataset(dataset, iterations, half, verbose)
+ .reshape([-1]) # type: ignore
+ )
 
  def predict_on_batch(self, data, iterations=1):
  """Rank the predictions according to their uncertainties."""

baal/active/heuristics/stochastics.py

@@ -7,7 +7,7 @@
 
 from baal.active.heuristics import AbstractHeuristic, Sequence
 
-log = structlog.get_logger(__name__)
+log = structlog.get_logger("baal")
 EPSILON = 1e-8
 
 

baal/bayesian/consistent_dropout.py

@@ -12,7 +12,7 @@
 
 class ConsistentDropout(_DropoutNd):
  """
- ConsistentDropout is useful when doing research.
+ ConsistentDropout is useful when doing learn.
  It guarantees that while the masks are the same between batches
  during inference. The masks are different inside the batch.
 
@@ -59,7 +59,7 @@ def eval(self):
 
 class ConsistentDropout2d(_DropoutNd):
  """
- ConsistentDropout is useful when doing research.
+ ConsistentDropout is useful when doing learn.
  It guarantees that while the mask are the same between batches,
  they are different inside the batch.
 

baal/calibration/calibration.py

@@ -11,7 +11,7 @@
 from baal.modelwrapper import TrainingArgs
 from baal.utils.metrics import ECE, ECE_PerCLs
 
-log = structlog.get_logger("Calibrating...")
+log = structlog.get_logger("baal")
 
 
 class DirichletCalibrator(object):

baal/experiments/__init__.py

@@ -0,0 +1,20 @@
+import abc
+from typing import Dict, List, Union
+
+from numpy._typing import NDArray
+
+from baal.active.dataset.base import Dataset
+
+
+class FrameworkAdapter(abc.ABC):
+ def reset_weights(self):
+ raise NotImplementedError
+
+ def train(self, al_dataset: Dataset) -> Dict[str, float]:
+ raise NotImplementedError
+
+ def predict(self, dataset: Dataset, iterations: int) -> Union[NDArray, List[NDArray]]:
+ raise NotImplementedError
+
+ def evaluate(self, dataset: Dataset, average_predictions: int) -> Dict[str, float]:
+ raise NotImplementedError

baal/experiments/base.py

@@ -0,0 +1,121 @@
+import itertools
+from typing import Union, Optional, Any
+
+import numpy as np
+import structlog
+from torch.utils.data import Subset
+from tqdm.autonotebook import tqdm
+
+from baal import ModelWrapper, ActiveLearningDataset
+from baal.active.dataset.base import Dataset
+from baal.active.heuristics import AbstractHeuristic
+from baal.active.stopping_criteria import StoppingCriterion, LabellingBudgetStoppingCriterion
+from baal.experiments import FrameworkAdapter
+from baal.experiments.modelwrapper import ModelWrapperAdapter
+
+try:
+ import transformers
+ from baal.transformers_trainer_wrapper import BaalTransformersTrainer
+ from baal.experiments.transformers import TransformersAdapter
+
+ TRANSFORMERS_AVAILABLE = True
+except ImportError:
+ BaalTransformersTrainer = None # type: ignore
+ TransformersAdapter = None # type: ignore
+ TRANSFORMERS_AVAILABLE = False
+
+log = structlog.get_logger("baal")
+
+
+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)

baal/experiments/modelwrapper.py

@@ -0,0 +1,28 @@
+from copy import deepcopy
+from typing import Dict, Union, List
+
+from numpy._typing import NDArray
+
+from baal import ModelWrapper
+from baal.active.dataset.base import Dataset
+from baal.experiments import FrameworkAdapter
+
+
+class ModelWrapperAdapter(FrameworkAdapter):
+ def __init__(self, wrapper: ModelWrapper):
+ self.wrapper = wrapper
+ self._init_weight = deepcopy(self.wrapper.state_dict())
+
+ def reset_weights(self):
+ self.wrapper.load_state_dict(self._init_weight)
+
+ def train(self, al_dataset: Dataset) -> Dict[str, float]:
+ self.wrapper.train_on_dataset(al_dataset)
+ return self.wrapper.get_metrics("train")
+
+ def predict(self, dataset: Dataset, iterations: int) -> Union[NDArray, List[NDArray]]:
+ return self.wrapper.predict_on_dataset(dataset, iterations=iterations)
+
+ def evaluate(self, dataset: Dataset, average_predictions: int) -> Dict[str, float]:
+ self.wrapper.test_on_dataset(dataset, average_predictions=average_predictions)
+ return self.wrapper.get_metrics("test")

baal/experiments/transformers.py

@@ -0,0 +1,29 @@
+from copy import deepcopy
+from typing import Dict, cast, List, Union
+
+from numpy._typing import NDArray
+
+from baal.active.dataset.base import Dataset
+from baal.experiments import FrameworkAdapter
+from baal.transformers_trainer_wrapper import BaalTransformersTrainer
+
+
+class TransformersAdapter(FrameworkAdapter):
+ def __init__(self, wrapper: BaalTransformersTrainer):
+ self.wrapper = wrapper
+ self._init_weight = deepcopy(self.wrapper.model.state_dict())
+
+ def reset_weights(self):
+ self.wrapper.model.load_state_dict(self._init_weight)
+ self.wrapper.lr_scheduler = None
+ self.wrapper.optimizer = None
+
+ def train(self, al_dataset: Dataset) -> Dict[str, float]:
+ self.wrapper.train_dataset = al_dataset
+ return self.wrapper.train().metrics # type: ignore
+
+ def predict(self, dataset: Dataset, iterations: int) -> Union[NDArray, List[NDArray]]:
+ return self.wrapper.predict_on_dataset(dataset, iterations=iterations)
+
+ def evaluate(self, dataset: Dataset, average_predictions: int) -> Dict[str, float]:
+ return cast(Dict[str, float], self.wrapper.evaluate(dataset))

baal/modelwrapper.py

@@ -3,15 +3,16 @@
 from collections.abc import Sequence
 from copy import deepcopy
 from dataclasses import dataclass
-from typing import Callable, Optional
+from typing import Callable, Optional, Union, List
 
 import numpy as np
 import structlog
 import torch
+from numpy._typing import NDArray
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader
 from torch.utils.data.dataloader import default_collate
-from tqdm import tqdm
+from tqdm.autonotebook import tqdm
 
 from baal.active.dataset.base import Dataset
 from baal.metrics.mixin import MetricMixin
@@ -22,7 +23,7 @@
 from baal.utils.metrics import Loss
 from baal.utils.warnings import raise_warnings_cache_replicated
 
-log = structlog.get_logger("ModelWrapper")
+log = structlog.get_logger("baal")
 
 
 def _stack_preds(out):
@@ -52,8 +53,7 @@ class ModelWrapper(MetricMixin):
 
  Args:
  model (nn.Module): The model to optimize.
- criterion (Callable): A loss function.
- replicate_in_memory (bool): Replicate in memory optional.
+ args (TrainingArgs): Model arguments for training/predicting.
  """
 
  def __init__(self, model, args: TrainingArgs):
@@ -148,7 +148,6 @@ def train_and_test_on_datasets(
  Args:
  train_dataset (Dataset): Dataset to train on.
  test_dataset (Dataset): Dataset to evaluate on.
- optimizer (Optimizer): Optimizer to use during 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.
@@ -236,7 +235,7 @@ def predict_on_dataset(
  iterations: int,
  half=False,
  verbose=True,
- ):
+ ) -> Union[NDArray, List[NDArray]]:
  """
  Use the model to predict on a dataset `iterations` time.