Repository: flairNLP/flair

Issue:
NER - Slow predictions on master branch compared to version 0.4.2
**Describe the bug**
With the very same code, and the very same data, Inference is much more slower on NER task when I use the version on master branch compared to the released version 0.4.2 (from 84 seconds to 234 seconds).
I have tried storing method None, cpu and gpu (btw documentation is missing, in particular none option is not clear to me even when reading the code)

FYI training the model doesn't take more time on master branch.
nvidia-smi tells me that during inference the GPU is used at 10-12% when using master, and at 25 to 30% when using 0.4.2

**To Reproduce**
```python
#  ...
    corpus: Corpus = prepare_flair_train_test_corpus(spacy_model=nlp, data_folder=data_folder, dev_size=dev_size)

    model_path = os.path.join(model_folder, 'best-model.pt')
    tagger: SequenceTagger = SequenceTagger.load(model_path)

    test_results, _ = tagger.evaluate(DataLoader(corpus.test, batch_size=50, num_workers=10))
    print(test_results.detailed_results)

    sentences_predict: List[Sentence] = corpus.train.sentences + corpus.test.sentences

    start = time.time()
    _ = tagger.predict(sentences_predict, mini_batch_size=50, verbose=True, embedding_storage_mode="gpu")
    print(time.time() - start)
    # 234 seconds on when using master branch
```
The training of the model looks like this:

```python
   # ...
    embedding_types: List[TokenEmbeddings] = [
        WordEmbeddings('fr'),
        FlairEmbeddings('resources/flair_ner/lm/ca_forward/best-lm.pt'),
        FlairEmbeddings('resources/flair_ner/lm/ca_backward/best-lm.pt'),
    ]

    embeddings: StackedEmbeddings = StackedEmbeddings(embeddings=embedding_types)

    tagger: SequenceTagger = SequenceTagger(hidden_size=256,
                                            embeddings=embeddings,
                                            tag_dictionary=tag_dictionary,
                                            tag_type='ner')
   # ...
```
Here it is a custom LM, but `fr` model from Zalendo has the same behaviour.

**Expected behavior**
I expect the time to process the same data with the same NER model to be the same than with master and 0.4.2.

**Environment (please complete the following information):**
Ubuntu 19.04
Flair master VS Flair 0.4.2
GPU 2080TI


Discussion:
Oha thanks for spotting this. Definitely something we need to investigate!
Again thanks for spotting this. I was able to reproduce and find the source of the problem. The new version does not only return the most likely tag but also a probability distribution over all tags in case you want to know what the second best guess was etc. However, it seems the implementation of this feature imposes a CPU bottleneck, hence the slowdown and reduced GPU usage. 

I'll push a PR that adds an optional flag to the `predict()` method that if not set will not do this distribution.` evaluate()` by default will not return the distribution. This speeds everything back up. Here the numbers of a v100: 

**Task - evaluating CoNLL-03 'ner' model**
_Release 0.4.2_:   ~30 seconds
_Current master_: ~40 seconds
_After fix_:             ~20 seconds

**Task - predict on all sentences in CoNLL-03 wirh 'ner' model**
_Release 0.4.2_:   ~100 seconds
_Current master_: ~190 seconds
_After fix_:             ~85 seconds



Resolving patch:
diff --git a/flair/models/sequence_tagger_model.py b/flair/models/sequence_tagger_model.py
--- a/flair/models/sequence_tagger_model.py
+++ b/flair/models/sequence_tagger_model.py
@@ -335,6 +335,7 @@ def predict(
         sentences: Union[List[Sentence], Sentence],
         mini_batch_size=32,
         embedding_storage_mode="none",
+        all_tag_prob: bool = False,
         verbose=False,
     ) -> List[Sentence]:
         with torch.no_grad():
@@ -366,14 +367,17 @@ def predict(
 
                 with torch.no_grad():
                     feature = self.forward(batch)
-                    tags, all_tags = self._obtain_labels(feature, batch)
+                    tags, all_tags = self._obtain_labels(
+                        feature, batch, get_all_tags=all_tag_prob
+                    )
 
                 for (sentence, sent_tags, sent_all_tags) in zip(batch, tags, all_tags):
                     for (token, tag, token_all_tags) in zip(
                         sentence.tokens, sent_tags, sent_all_tags
                     ):
                         token.add_tag_label(self.tag_type, tag)
-                        token.add_tags_proba_dist(self.tag_type, token_all_tags)
+                        if all_tag_prob:
+                            token.add_tags_proba_dist(self.tag_type, token_all_tags)
 
                 # clearing token embeddings to save memory
                 store_embeddings(batch, storage_mode=embedding_storage_mode)
@@ -530,7 +534,7 @@ def _calculate_loss(
             return score
 
     def _obtain_labels(
-        self, feature, sentences
+        self, feature, sentences, get_all_tags: bool = False
     ) -> (List[List[Label]], List[List[List[Label]]]):
         """
         Returns a tuple of two lists:
@@ -545,7 +549,9 @@ def _obtain_labels(
         all_tags = []
         for feats, length in zip(feature, lengths):
             if self.use_crf:
-                confidences, tag_seq, scores = self._viterbi_decode(feats[:length])
+                confidences, tag_seq, scores = self._viterbi_decode(
+                    feats[:length], all_scores=get_all_tags
+                )
             else:
                 tag_seq = []
                 confidences = []
@@ -565,22 +571,24 @@ def _obtain_labels(
                 ]
             )
 
-            all_tags.append(
-                [
+            if get_all_tags:
+                all_tags.append(
                     [
-                        Label(self.tag_dictionary.get_item_for_index(score_id), score)
-                        for score_id, score in enumerate(score_dist)
+                        [
+                            Label(
+                                self.tag_dictionary.get_item_for_index(score_id), score
+                            )
+                            for score_id, score in enumerate(score_dist)
+                        ]
+                        for score_dist in scores
                     ]
-                    for score_dist in scores
-                ]
-            )
+                )
 
         return tags, all_tags
 
-    def _viterbi_decode(self, feats):
+    def _viterbi_decode(self, feats, all_scores: bool = False):
         backpointers = []
         backscores = []
-        scores = []
 
         init_vvars = (
             torch.FloatTensor(1, self.tagset_size).to(flair.device).fill_(-10000.0)
@@ -621,25 +629,31 @@ def _viterbi_decode(self, feats):
             _, idx = torch.max(backscore, 0)
             prediction = idx.item()
             best_scores.append(softmax[prediction].item())
-            scores.append([elem.item() for elem in softmax.flatten()])
 
         start = best_path.pop()
         assert start == self.tag_dictionary.get_idx_for_item(START_TAG)
         best_path.reverse()
 
-        for index, (tag_id, tag_scores) in enumerate(zip(best_path, scores)):
-            if type(tag_id) != int and tag_id.item() != np.argmax(tag_scores):
-                swap_index_score = np.argmax(tag_scores)
-                scores[index][tag_id.item()], scores[index][swap_index_score] = (
-                    scores[index][swap_index_score],
-                    scores[index][tag_id.item()],
-                )
-            elif type(tag_id) == int and tag_id != np.argmax(tag_scores):
-                swap_index_score = np.argmax(tag_scores)
-                scores[index][tag_id], scores[index][swap_index_score] = (
-                    scores[index][swap_index_score],
-                    scores[index][tag_id],
-                )
+        scores = []
+        # return all scores if so selected
+        if all_scores:
+            for backscore in backscores:
+                softmax = F.softmax(backscore, dim=0)
+                scores.append([elem.item() for elem in softmax.flatten()])
+
+            for index, (tag_id, tag_scores) in enumerate(zip(best_path, scores)):
+                if type(tag_id) != int and tag_id.item() != np.argmax(tag_scores):
+                    swap_index_score = np.argmax(tag_scores)
+                    scores[index][tag_id.item()], scores[index][swap_index_score] = (
+                        scores[index][swap_index_score],
+                        scores[index][tag_id.item()],
+                    )
+                elif type(tag_id) == int and tag_id != np.argmax(tag_scores):
+                    swap_index_score = np.argmax(tag_scores)
+                    scores[index][tag_id], scores[index][swap_index_score] = (
+                        scores[index][swap_index_score],
+                        scores[index][tag_id],
+                    )
 
         return best_scores, best_path, scores
 


Test patch: