Repository: flairNLP/flair

Issue:
Fix cuda out of memory issue for text classification
When training a text classifier on large datasets, the trainer runs into memory problems. 
Investigate why we are running into out of memory issues and fix them.


Discussion:
I can confirm that issue. Currently I'm trying to train on the GermEval 2018 dataset, and passing `embeddings_in_memory=False` to the `train` function helps a bit. I also removed evaluation on training dataset and then no "cuda out of memory" messages appear.

**Update**: results are not very promising at the moment.
@stefan-it It seems like the evaluation method takes a lot of memory and causes the `cuda out of memory` problems. We are looking into it.

Yesterday we fixed a bug in training a text classifier on a single label task (see https://github.com/zalandoresearch/flair/issues/109). So in case you used single label training, you might want to try again with the current version in the master branch.

Resolving patch:
diff --git a/flair/embeddings.py b/flair/embeddings.py
--- a/flair/embeddings.py
+++ b/flair/embeddings.py
@@ -307,8 +307,6 @@ class CharLMEmbeddings(TokenEmbeddings):
     """Contextual string embeddings of words, as proposed in Akbik et al., 2018."""
 
     def __init__(self, model, detach: bool = True):
-        super().__init__()
-
         """
             Contextual string embeddings of words, as proposed in Akbik et al., 2018.
 
@@ -321,6 +319,7 @@ def __init__(self, model, detach: bool = True):
                 if set to false, the gradient will propagate into the language model. this dramatically slows down
                 training and often leads to worse results, so not recommended.
         """
+        super().__init__()
 
         # news-english-forward
         if model.lower() == 'news-forward':
@@ -381,7 +380,6 @@ def embedding_length(self) -> int:
         return self.__embedding_length
 
     def _add_embeddings_internal(self, sentences: List[Sentence]) -> List[Sentence]:
-
         # get text sentences
         text_sentences = [sentence.to_tokenized_string() for sentence in sentences]
 
diff --git a/flair/trainers/text_classification_trainer.py b/flair/trainers/text_classification_trainer.py
--- a/flair/trainers/text_classification_trainer.py
+++ b/flair/trainers/text_classification_trainer.py
@@ -29,21 +29,26 @@ def train(self,
               base_path: str,
               learning_rate: float = 0.1,
               mini_batch_size: int = 32,
+              eval_mini_batch_size: int = 8,
               max_epochs: int = 100,
               anneal_factor: float = 0.5,
               patience: int = 2,
               save_model: bool = True,
               embeddings_in_memory: bool = True,
-              train_with_dev: bool = False):
+              train_with_dev: bool = False,
+              eval_on_train: bool = False):
         """
         Trains the model using the training data of the corpus.
         :param base_path: the directory to which any results should be written to
         :param learning_rate: the learning rate
         :param mini_batch_size: the mini batch size
+        :param eval_mini_batch_size: the mini batch size for evaluation
         :param max_epochs: the maximum number of epochs to train
         :param save_model: boolean value indicating, whether the model should be saved or not
         :param embeddings_in_memory: boolean value indicating, if embeddings should be kept in memory or not
         :param train_with_dev: boolean value indicating, if the dev data set should be used for training or not
+        :param eval_on_train: boolean value indicating, if evaluation metrics should be calculated on training data set
+        or not
         """
 
         loss_txt = init_output_file(base_path, 'loss.txt')
@@ -100,8 +105,11 @@ def train(self,
                         clear_embeddings(batch)
 
                     if batch_no % modulo == 0:
-                        print("epoch {0} - iter {1}/{2} - loss {3:.8f}".format(epoch + 1, batch_no, len(batches),
-                                                                               current_loss / seen_sentences))
+                        for group in optimizer.param_groups:
+                            learning_rate = group['lr']
+
+                        print("epoch {0} - iter {1}/{2} - loss {3:.8f} - lr {4:.4f} - bad epochs {5}".format(
+                            epoch + 1, batch_no, len(batches), current_loss / seen_sentences, learning_rate, scheduler.num_bad_epochs))
                         iteration = epoch * len(batches) + batch_no
                         self._extract_weights(iteration, weights_index, weights_txt)
 
@@ -111,13 +119,15 @@ def train(self,
 
                 print('-' * 100)
 
-                train_acc, train_f_score, train_loss = self._calculate_evaluation_results_for(
-                    'TRAIN', self.corpus.train, embeddings_in_memory, epoch, mini_batch_size)
+                train_f_score = train_acc = train_loss = 0
+                if eval_on_train:
+                    train_acc, train_f_score, train_loss = self._calculate_evaluation_results_for(
+                        'TRAIN', self.corpus.train, embeddings_in_memory, epoch, eval_mini_batch_size)
 
                 dev_f_score = dev_acc = dev_loss = 0
                 if not train_with_dev:
                     dev_acc, dev_f_score, dev_loss = self._calculate_evaluation_results_for(
-                        'DEV', self.corpus.dev, embeddings_in_memory, epoch, mini_batch_size)
+                        'DEV', self.corpus.dev, embeddings_in_memory, epoch, eval_mini_batch_size)
 
                 with open(loss_txt, 'a') as f:
                     f.write('{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\n'.format(
@@ -179,7 +189,7 @@ def _calculate_evaluation_results_for(self, dataset_name, dataset, embeddings_in
 
         return acc, f_score, loss
 
-    def evaluate(self, sentences: List[Sentence], eval_class_metrics: bool = False, mini_batch_size: int = 32,
+    def evaluate(self, sentences: List[Sentence], eval_class_metrics: bool = False, mini_batch_size: int = 16,
                  embeddings_in_memory: bool = True) -> (dict, float):
         """
         Evaluates the model with the given list of sentences.
@@ -200,13 +210,13 @@ def evaluate(self, sentences: List[Sentence], eval_class_metrics: bool = False,
             labels = self.model.obtain_labels(scores)
             loss = self.model.calculate_loss(scores, batch)
 
+            if not embeddings_in_memory:
+                clear_embeddings(batch)
+
             eval_loss += loss
 
             y_pred.extend(convert_labels_to_one_hot([[label.name for label in sent_labels] for sent_labels in labels], self.label_dict))
 
-            if not embeddings_in_memory:
-                clear_embeddings(batch)
-
         metrics = [calculate_micro_avg_metric(y_true, y_pred, self.label_dict)]
         if eval_class_metrics:
             metrics.extend(calculate_class_metrics(y_true, y_pred, self.label_dict))


Test patch: