Repository: lightly-ai/lightly

Issue:
Adapt active learning score to predict for all and upload to query tag
## Description
Following our discussions around #482, we need to change two things to implement B1:
- [x] upload the scores to the query tag
- [x] compute the scores for the whole query tag, not only the unlabeled set
## Tasks
- [x] make the changes in the description
- [x] update the tutorials
- [x] update the documentation

## Acceptance Criteria
- [x] active learning tutorial still works


Discussion:

Resolving patch:
diff --git a/docs/source/tutorials_source/platform/tutorial_active_learning.py b/docs/source/tutorials_source/platform/tutorial_active_learning.py
--- a/docs/source/tutorials_source/platform/tutorial_active_learning.py
+++ b/docs/source/tutorials_source/platform/tutorial_active_learning.py
@@ -188,9 +188,9 @@ def get_labels(self, filenames: List[str]) -> np.ndarray:
 classifier.fit(X=labeled_set_features, y=labeled_set_labels)
 
 # %%
-# 3. Use the classifier to predict on the unlabeled set.
-unlabeled_set_features = dataset.get_features(agent.unlabeled_set)
-predictions = classifier.predict_proba(X=unlabeled_set_features)
+# 3. Use the classifier to predict on the query set.
+query_set_features = dataset.get_features(agent.query_set)
+predictions = classifier.predict_proba(X=query_set_features)
 
 # %%
 # 4. Calculate active learning scores from the prediction.
diff --git a/lightly/active_learning/agents/agent.py b/lightly/active_learning/agents/agent.py
--- a/lightly/active_learning/agents/agent.py
+++ b/lightly/active_learning/agents/agent.py
@@ -14,14 +14,18 @@ class ActiveLearningAgent:
     Attributes:
         api_workflow_client:
             The client to connect to the api.
-        preselected_tag_id:
-            The id of the tag containing the already labeled samples, default: None == no labeled samples yet.
-        query_tag_id:
-            The id of the tag defining where to sample from, default: None resolves to initial_tag
+        query_set:
+            Set of filenames corresponding to samples which can possibly be selected.
+            Set to all samples in the query tag or to the whole dataset by default.
         labeled_set:
-            The filenames of the samples in the labeled set, List[str]
+            Set of filenames corresponding to samples in the labeled set.
+            Set to all samples in the preselected tag or to an empty list by default.
         unlabeled_set:
-            The filenames of the samples in the unlabeled set, List[str]
+            Set of filenames corresponding to samples which are in the query set
+            but not in the labeled set.
+        added_set:
+            Set of filenames corresponding to samples which were added to the 
+            labeled set in the last query.
 
     Examples:
         >>> # set the token and dataset id
@@ -34,78 +38,137 @@ class ActiveLearningAgent:
         >>>
         >>> # make an initial active learning query
         >>> sampler_config = SamplerConfig(n_samples=100, name='initial-set')
-        >>> initial_set = agent.query(sampler_config)
-        >>> unlabeled_set = agent.unlabeled_set
+        >>> agent.query(sampler_config)
+        >>> initial_set = agent.labeled_set
         >>>
         >>> # train and evaluate a model on the initial set
-        >>> # make predictions on the unlabeled set (keep ordering of filenames)
+        >>> # make predictions on the query set:
+        >>> query_set = agent.query_set
+        >>> # important:
+        >>> # be sure to keep the order of the query set when you make predictions
         >>>
         >>> # create active learning scorer
         >>> scorer = ScorerClassification(predictions)
         >>>
         >>> # make a second active learning query
         >>> sampler_config = SamplerConfig(n_samples=200, name='second-set')
-        >>> second_set = agent.query(sampler_config, scorer)
+        >>> agent.query(sampler_config, scorer)
+        >>> added_set = agent.added_set # access only the samples added by this query
 
     """
 
-    def __init__(self, api_workflow_client: ApiWorkflowClient, query_tag_name: str = None,
+    def __init__(self,
+                 api_workflow_client: ApiWorkflowClient,
+                 query_tag_name: str = 'initial-tag',
                  preselected_tag_name: str = None):
 
         self.api_workflow_client = api_workflow_client
-        if query_tag_name is not None or preselected_tag_name is not None:
-            tag_name_id_dict = dict([tag.name, tag.id] for tag in self.api_workflow_client._get_all_tags())
-            if preselected_tag_name is not None:
-                self.preselected_tag_id = tag_name_id_dict[preselected_tag_name]
-            if query_tag_name is not None:
-                self.query_tag_id = tag_name_id_dict[query_tag_name]
 
-        if not hasattr(self, "preselected_tag_id"):
-            self.preselected_tag_id = None
-        if not hasattr(self, "query_tag_id"):
-            self.query_tag_id = None
-        self._set_labeled_and_unlabeled_set()
+        # set the query_tag_id and preselected_tag_id
+        self._query_tag_id = None
+        self._preselected_tag_id = None
 
-    def _set_labeled_and_unlabeled_set(self, preselected_tag_data: TagData = None):
-        """Sets the labeled and unlabeled set based on the preselected and query tag id
+        # build lookup table for tag_name to tag_id
+        tag_name_id_dict = {}
+        for tag in self.api_workflow_client._get_all_tags():
+            tag_name_id_dict[tag.name] = tag.id
+        # use lookup table to set ids
+        self._query_tag_id = tag_name_id_dict[query_tag_name]
+        if preselected_tag_name is not None:
+            self._preselected_tag_id = tag_name_id_dict[preselected_tag_name]
 
-        It loads the bitmaks for the both tag_ids from the server and then
-        extracts the filenames from it given the mapping on the server.
+        # set the filename lists based on preselected and query tag
+        self._query_tag_bitmask = self._get_query_tag_bitmask()
+        self._preselected_tag_bitmask = self._get_preselected_tag_bitmask()
+        # keep track of the last preselected tag to compute added samples
+        self._old_preselected_tag_bitmask = None
 
-        Args:
-            preselected_tag_data:
-                optional param, then it must not be loaded from the API
+
+    def _get_query_tag_bitmask(self):
+        """Initializes the query tag bitmask.
 
         """
+        # get query tag from api and set bitmask accordingly
+        query_tag_data = self.api_workflow_client.tags_api.get_tag_by_tag_id(
+            self.api_workflow_client.dataset_id,
+            tag_id=self._query_tag_id
+        )
+        query_tag_bitmask = BitMask.from_hex(query_tag_data.bit_mask_data)
+
+        return query_tag_bitmask
 
-        if not hasattr(self, "bitmask_labeled_set"):
-            self.bitmask_labeled_set = BitMask.from_hex("0x0")  # empty labeled set
-            self.bitmask_added_set = BitMask.from_hex("0x0")  # empty added set
-        if self.preselected_tag_id is not None:  # else the default values (empty labeled and added set) are kept
-            if preselected_tag_data is None:  # if it is not passed as argument, it must be loaded from the API
-                preselected_tag_data = self.api_workflow_client.tags_api.get_tag_by_tag_id(
-                    self.api_workflow_client.dataset_id, tag_id=self.preselected_tag_id)
-            new_bitmask_labeled_set = BitMask.from_hex(preselected_tag_data.bit_mask_data)
-            self.bitmask_added_set = new_bitmask_labeled_set - self.bitmask_labeled_set
-            self.bitmask_labeled_set = new_bitmask_labeled_set
-
-        if self.query_tag_id is None:
-            bitmask_query_tag = BitMask.from_length(len(self.api_workflow_client.filenames_on_server))
+    def _get_preselected_tag_bitmask(self):
+        """Initializes the preselected tag bitmask.
+
+        """
+        if self._preselected_tag_id is None:
+            # if not specified, no samples belong to the preselected tag
+            preselected_tag_bitmask = BitMask.from_hex('0x0')
         else:
-            query_tag_data = self.api_workflow_client.tags_api.get_tag_by_tag_id(
-                self.api_workflow_client.dataset_id, tag_id=self.query_tag_id)
-            bitmask_query_tag = BitMask.from_hex(query_tag_data.bit_mask_data)
-        self.bitmask_unlabeled_set = bitmask_query_tag - self.bitmask_labeled_set
+            # get preselected tag from api and set bitmask accordingly
+            preselected_tag_data = self.api_workflow_client.tags_api.get_tag_by_tag_id(
+                self.api_workflow_client.dataset_id,
+                tag_id=self._preselected_tag_id
+            )
+            preselected_tag_bitmask = BitMask.from_hex(preselected_tag_data.bit_mask_data)
+
+        return preselected_tag_bitmask
+
+    @property
+    def query_set(self):
+        """List of filenames for which to calculate active learning scores.
+
+        """
+        return self._query_tag_bitmask.masked_select_from_list(
+            self.api_workflow_client.filenames_on_server
+        )
+
+    @property
+    def labeled_set(self):
+        """List of filenames indicating selected samples.
+
+        """
+        return self._preselected_tag_bitmask.masked_select_from_list(
+            self.api_workflow_client.filenames_on_server
+        )
+
+    @property
+    def unlabeled_set(self):
+        """List of filenames which belong to the query set but are not selected.
+
+        """
+        # unlabeled set is the query set minus the preselected set
+        unlabeled_tag_bitmask = self._query_tag_bitmask - self._preselected_tag_bitmask
+        return unlabeled_tag_bitmask.masked_select_from_list(
+            self.api_workflow_client.filenames_on_server
+        )
+
+    @property
+    def added_set(self):
+        """List of filenames of newly added samples (in the last query).
 
-        self.labeled_set = self.bitmask_labeled_set.masked_select_from_list(self.api_workflow_client.filenames_on_server)
-        self.added_set = self.bitmask_added_set.masked_select_from_list(self.api_workflow_client.filenames_on_server)
-        self.unlabeled_set = self.bitmask_unlabeled_set.masked_select_from_list(self.api_workflow_client.filenames_on_server)
+        Raises:
+            RuntimeError if executed before a query.
 
-    def query(self, sampler_config: SamplerConfig, al_scorer: Scorer = None) -> Tuple[List[str], List[str]]:
+        """
+        # the added set only exists after a query
+        if self._old_preselected_tag_bitmask is None:
+            raise RuntimeError('Cannot compute \"added set\" before querying.')
+        # added set is new preselected set minus the old one
+        added_tag_bitmask = self._preselected_tag_bitmask - self._old_preselected_tag_bitmask
+        return added_tag_bitmask.masked_select_from_list(
+            self.api_workflow_client.filenames_on_server
+        )
+
+
+    def query(self,
+              sampler_config: SamplerConfig,
+              al_scorer: Scorer = None) -> Tuple[List[str], List[str]]:
         """Performs an active learning query.
 
-        As part of it, the self.labeled_set and self.unlabeled_set are updated
-        and can be used for the next step.
+        After the query, the labeled set is updated to contain all selected samples,
+        the added set is recalculated as (new labeled set - old labeled set), and
+        the query set stays the same.
 
         Args:
             sampler_config:
@@ -113,43 +176,39 @@ def query(self, sampler_config: SamplerConfig, al_scorer: Scorer = None) -> Tupl
             al_scorer:
                 An instance of a class inheriting from Scorer, e.g. a ClassificationScorer.
 
-        Returns:
-            The filenames of the samples in the new labeled_set
-            and the filenames of the samples chosen by the sampler.
-            This added_set was added to the old labeled_set
-            to form the new labeled_set.
-
         """
-        # check input
+
+        # handle illogical stopping condition
         if sampler_config.n_samples < len(self.labeled_set):
-            warnings.warn("ActiveLearningAgent.query: The number of samples which should be sampled "
-                          "including the current labeled set "
-                          "(sampler_config.n_samples) "
-                          "is smaller than the number of samples in the current labeled set."
-                          "Skipping the sampling and returning the previous labeled set.")
-            return self.labeled_set, []
-
-        # calculate scores
+            warnings.warn(
+                f'ActiveLearningAgent.query: The number of samples ({sampler_config.n_samples}) is '
+                f'smaller than the number of preselected samples ({len(self.labeled_set)}).'
+                'Skipping the active learning query.'
+            )
+            return
+
+        # calculate active learning scores
+        scores_dict = None
         if al_scorer is not None:
-            no_unlabeled_samples = len(self.unlabeled_set)
-            no_samples_with_predictions = len(al_scorer.model_output)
-            if no_unlabeled_samples != no_samples_with_predictions:
-                raise ValueError(f"The scorer must have exactly as many samples as in the unlabeled set,"
-                                 f"but there are {no_samples_with_predictions} predictions in the scorer,"
-                                 f"but {no_unlabeled_samples} in the unlabeled set.")
+            no_query_samples = len(self.query_set)
+            no_query_samples_with_scores = len(al_scorer.model_output)
+            if no_query_samples != no_query_samples_with_scores:
+                raise ValueError(
+                    f'Number of query samples ({no_query_samples}) must match '
+                    f'the number of predictions ({no_query_samples_with_scores})!'
+                )
             scores_dict = al_scorer.calculate_scores()
-        else:
-            scores_dict = None
 
         # perform the sampling
         new_tag_data = self.api_workflow_client.sampling(
             sampler_config=sampler_config,
             al_scores=scores_dict,
-            preselected_tag_id=self.preselected_tag_id,
-            query_tag_id=self.query_tag_id)
-
-        # set the newly chosen tag as the new preselected_tag_id and update the sets
-        self.preselected_tag_id = new_tag_data.id
-        self._set_labeled_and_unlabeled_set(new_tag_data)
-
-        return self.labeled_set, self.added_set
+            preselected_tag_id=self._preselected_tag_id,
+            query_tag_id=self._query_tag_id
+        )
+
+        # update the old preselected_tag
+        self._old_preselected_tag_bitmask = self._preselected_tag_bitmask
+        # set the newly chosen tag as the new preselected_tag
+        self._preselected_tag_id = new_tag_data.id
+        self._preselected_tag_bitmask = self._get_preselected_tag_bitmask()
diff --git a/lightly/api/api_workflow_sampling.py b/lightly/api/api_workflow_sampling.py
--- a/lightly/api/api_workflow_sampling.py
+++ b/lightly/api/api_workflow_sampling.py
@@ -10,8 +10,21 @@
 from lightly.openapi_generated.swagger_client.models.tag_data import TagData
 from lightly.openapi_generated.swagger_client.models.sampling_config import SamplingConfig
 from lightly.openapi_generated.swagger_client.models.sampling_create_request import SamplingCreateRequest
-from lightly.openapi_generated.swagger_client.models.sampling_config_stopping_condition import \
-    SamplingConfigStoppingCondition
+from lightly.openapi_generated.swagger_client.models.sampling_config_stopping_condition import SamplingConfigStoppingCondition
+
+
+
+def _parse_active_learning_scores(scores):
+    """Makes list/np.array of active learning scores serializable.
+
+    """
+    # the api only accepts float64s
+    if isinstance(scores, np.ndarray):
+        scores = scores.astype(np.float64)
+
+    # convert to list and return
+    return list(scores)
+
 
 
 class _SamplingMixin:
@@ -41,7 +54,9 @@ def sampling(self, sampler_config: SamplerConfig, al_scores: Dict[str, List[np.n
         # make sure the tag name does not exist yet
         tags = self._get_all_tags()
         if sampler_config.name in [tag.name for tag in tags]:
-            raise RuntimeError(f"There already exists a tag with tag_name {sampler_config.name}.")
+            raise RuntimeError(f'There already exists a tag with tag_name {sampler_config.name}.')
+        if len(tags) == 0:
+            raise RuntimeError('There exists no initial-tag for this dataset.')
 
         # make sure we have an embedding id
         try:
@@ -51,15 +66,25 @@ def sampling(self, sampler_config: SamplerConfig, al_scores: Dict[str, List[np.n
 
 
         # upload the active learning scores to the api
+        # change @20210422: we store the active learning scores with the query
+        # tag. policy is that if there's no explicit query tag, the whole dataset
+        # will be the query tag (i.e. query_tag = initial-tag)
         if al_scores is not None:
-            if preselected_tag_id is None:
-                raise ValueError
+            # set the query tag to the initial-tag if necessary
+            if query_tag_id is None:
+                query_tag = next(t for t in tags if t.name == 'initial-tag')
+                query_tag_id = query_tag.id
+            # iterate over all available score types and upload them
             for score_type, score_values in al_scores.items():
-                if isinstance(score_values, np.ndarray):
-                    score_values = score_values.astype(np.float64)
-                body = ActiveLearningScoreCreateRequest(score_type=score_type, scores=list(score_values))
+                body = ActiveLearningScoreCreateRequest(
+                    score_type=score_type,
+                    scores=_parse_active_learning_scores(score_values)
+                )
                 self.scores_api.create_or_update_active_learning_score_by_tag_id(
-                    body, dataset_id=self.dataset_id, tag_id=preselected_tag_id)
+                    body,
+                    dataset_id=self.dataset_id,
+                    tag_id=query_tag_id,
+                )
 
         # trigger the sampling
         payload = self._create_sampling_create_request(sampler_config, preselected_tag_id, query_tag_id)


Test patch:
diff --git a/tests/UNMOCKED_end2end_tests/test_api.py b/tests/UNMOCKED_end2end_tests/test_api.py
--- a/tests/UNMOCKED_end2end_tests/test_api.py
+++ b/tests/UNMOCKED_end2end_tests/test_api.py
@@ -100,7 +100,7 @@ def create_new_dataset_with_embeddings(path_to_dataset: str,
 
 def t_est_active_learning(api_workflow_client: ApiWorkflowClient,
                           method: SamplingMethod = SamplingMethod.CORAL,
-                          query_tag_name: str = None,
+                          query_tag_name: str = 'initial-tag',
                           preselected_tag_name: str = None,
                           n_samples_additional: List[int] = [2, 5]):
     # create the tags with 100 respectively 10 samples if not yet existant
@@ -142,7 +142,7 @@ def t_est_active_learning(api_workflow_client: ApiWorkflowClient,
         assert len(agent.unlabeled_set) == total_no_samples - n_samples
 
         # Update the scorer
-        n_samples = len(agent.unlabeled_set)
+        n_samples = len(agent.query_set)
         n_classes = 10
         predictions = np.random.rand(n_samples, n_classes)
         predictions_normalized = predictions / np.sum(predictions, axis=1)[:, np.newaxis]
@@ -164,7 +164,7 @@ def t_est_api_with_matrix(path_to_dataset: str,
     )
 
     for method in [SamplingMethod.CORAL, SamplingMethod.CORESET, SamplingMethod.RANDOM]:
-        for query_tag_name in [None, "query_tag_name_xyz"]:
+        for query_tag_name in ['initial-tag', "query_tag_name_xyz"]:
             for preselected_tag_name in [None, "preselected_tag_name_xyz"]:
                 print(f"Starting AL run with method '{method}', query_tag '{query_tag_name}' "
                       f"and preselected_tag '{preselected_tag_name}'.")
diff --git a/tests/active_learning/test_active_learning_agent.py b/tests/active_learning/test_active_learning_agent.py
--- a/tests/active_learning/test_active_learning_agent.py
+++ b/tests/active_learning/test_active_learning_agent.py
@@ -30,31 +30,63 @@ def test_agent(self):
                         sampler_config = SamplerConfig(n_samples=n_samples, method=method)
 
                     if sampler_config.method == SamplingMethod.CORAL:
-                        predictions = np.random.rand(len(agent.unlabeled_set), 10).astype(np.float32)
+                        predictions = np.random.rand(len(agent.query_set), 10).astype(np.float32)
                         predictions_normalized = predictions / np.sum(predictions, axis=1)[:, np.newaxis]
                         al_scorer = ScorerClassification(predictions_normalized)
-                        labeled_set, added_set = agent.query(sampler_config=sampler_config, al_scorer=al_scorer)
+                        agent.query(sampler_config=sampler_config, al_scorer=al_scorer)
                     else:
                         sampler_config = SamplerConfig(n_samples=n_samples)
-                        labeled_set, added_set = agent.query(sampler_config=sampler_config)
+                        agent.query(sampler_config=sampler_config)
+                    
+                    labeled_set, added_set = agent.labeled_set, agent.added_set
 
                     self.assertEqual(n_old_labeled + len(added_set), len(labeled_set))
-                    assert set(added_set).issubset(labeled_set)
+                    self.assertTrue(set(added_set).issubset(labeled_set))
                     self.assertEqual(len(list(set(agent.labeled_set) & set(agent.unlabeled_set))), 0)
                     self.assertEqual(n_old_unlabeled - len(added_set), len(agent.unlabeled_set))
 
-    def test_agent_wrong_scores(self):
+    def test_agent_wrong_number_of_scores(self):
         self.api_workflow_client.embedding_id = "embedding_id_xyz"
 
         agent = ActiveLearningAgent(self.api_workflow_client, preselected_tag_name="preselected_tag_name_xyz")
         method = SamplingMethod.CORAL
         n_samples = len(agent.labeled_set) + 2
 
-        n_predictions = len(agent.unlabeled_set) - 3  # the -3 should cause en error
+        n_predictions = len(agent.query_set) - 3 # the -3 should cause en error
         predictions = np.random.rand(n_predictions, 10).astype(np.float32)
         predictions_normalized = predictions / np.sum(predictions, axis=1)[:, np.newaxis]
         al_scorer = ScorerClassification(predictions_normalized)
 
         sampler_config = SamplerConfig(n_samples=n_samples, method=method)
         with self.assertRaises(ValueError):
-            labeled_set, added_set = agent.query(sampler_config=sampler_config, al_scorer=al_scorer)
+            agent.query(sampler_config=sampler_config, al_scorer=al_scorer)
+
+    def test_agent_added_set_before_query(self):
+
+        self.api_workflow_client.embedding_id = "embedding_id_xyz"
+        agent = ActiveLearningAgent(
+            self.api_workflow_client,
+            preselected_tag_name="preselected_tag_name_xyz"
+        )
+
+        agent.query_set
+        agent.labeled_set
+        agent.unlabeled_set
+        with self.assertRaises(RuntimeError):
+            agent.added_set
+
+    def test_agent_query_too_few(self):
+
+        self.api_workflow_client.embedding_id = "embedding_id_xyz"
+        agent = ActiveLearningAgent(
+            self.api_workflow_client,
+            preselected_tag_name="preselected_tag_name_xyz",
+        )
+
+        # sample 0 samples
+        sampler_config = SamplerConfig(
+            n_samples=0,
+            method=SamplingMethod.RANDOM
+        )
+
+        agent.query(sampler_config)