Repository: flairNLP/flair

Issue:
ModelTrainer annealing options
Currently, the `train_with_dev` parameter in the `ModelTrainer` impacts the annealing behavior: If set to True, we anneal against training loss, otherwise against current score computed on dev data. 

This is a bit problematic since: 

- We may want more control over annealing behavior: For instance, we might want to anneal against train loss even if train_with_dev is set to False.
- This behavior is implicit, while is may be easier to understand if it can be explicitly set 

**Task**: Add more annealing options to `ModelTrainer` to make it possible to explicitly control annealing behavior.


Discussion:

Resolving patch:
diff --git a/flair/embeddings.py b/flair/embeddings.py
--- a/flair/embeddings.py
+++ b/flair/embeddings.py
@@ -6,7 +6,6 @@
 import gensim
 import numpy as np
 import torch
-from torch.utils.data import TensorDataset, DataLoader, SequentialSampler
 from deprecated import deprecated
 
 from pytorch_pretrained_bert.tokenization import BertTokenizer
diff --git a/flair/hyperparameter/param_selection.py b/flair/hyperparameter/param_selection.py
--- a/flair/hyperparameter/param_selection.py
+++ b/flair/hyperparameter/param_selection.py
@@ -96,9 +96,9 @@ def _objective(self, params: dict):
 
             # take the average over the last three scores of training
             if self.optimization_value == OptimizationValue.DEV_LOSS:
-                curr_scores = result['loss_history'][-3:]
+                curr_scores = result['dev_loss_history'][-3:]
             else:
-                curr_scores = list(map(lambda s: 1 - s, result['score_history'][-3:]))
+                curr_scores = list(map(lambda s: 1 - s, result['dev_score_history'][-3:]))
             score = sum(curr_scores) / float(len(curr_scores))
             var = np.var(curr_scores)
             scores.append(score)
diff --git a/flair/trainers/trainer.py b/flair/trainers/trainer.py
--- a/flair/trainers/trainer.py
+++ b/flair/trainers/trainer.py
@@ -1,3 +1,4 @@
+from enum import Enum
 from pathlib import Path
 from typing import List, Union
 
@@ -35,78 +36,7 @@ def __init__(self,
         self.loss: float = loss
         self.scheduler_state: dict = scheduler_state
         self.optimizer_state: dict = optimizer_state
-
-    def find_learning_rate(self,
-                           base_path: Union[Path, str],
-                           file_name: str = 'learning_rate.tsv',
-                           start_learning_rate: float = 1e-7,
-                           end_learning_rate: float = 10,
-                           iterations: int = 100,
-                           mini_batch_size: int = 32,
-                           stop_early: bool = True,
-                           smoothing_factor: float = 0.98,
-                           **kwargs
-                           ) -> Path:
-        best_loss = None
-        moving_avg_loss = 0
-
-        # cast string to Path
-        if type(base_path) is str:
-            base_path = Path(base_path)
-        learning_rate_tsv = init_output_file(base_path, file_name)
-
-        with open(learning_rate_tsv, 'a') as f:
-            f.write('ITERATION\tTIMESTAMP\tLEARNING_RATE\tTRAIN_LOSS\n')
-
-        optimizer = self.optimizer(self.model.parameters(), lr=start_learning_rate, **kwargs)
-
-        train_data = self.corpus.train
-        random.shuffle(train_data)
-        batches = [train_data[x:x + mini_batch_size] for x in range(0, len(train_data), mini_batch_size)][:iterations]
-
-        scheduler = ExpAnnealLR(optimizer, end_learning_rate, iterations)
-
-        model_state = self.model.state_dict()
-        model_device = next(self.model.parameters()).device
-        self.model.train()
-
-        for itr, batch in enumerate(batches):
-            loss = self.model.forward_loss(batch)
-
-            optimizer.zero_grad()
-            loss.backward()
-            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5.0)
-            optimizer.step()
-            scheduler.step()
-            learning_rate = scheduler.get_lr()[0]
-
-            loss_item = loss.item()
-            if itr == 0:
-                best_loss = loss_item
-            else:
-                if smoothing_factor > 0:
-                    moving_avg_loss = smoothing_factor * moving_avg_loss + (1 - smoothing_factor) * loss_item
-                    loss_item = moving_avg_loss / (1 - smoothing_factor ** (itr+1))
-                if loss_item < best_loss:
-                    best_loss = loss
-
-            if stop_early and (loss_item > 4 * best_loss or torch.isnan(loss)):
-                log_line(log)
-                log.info('loss diverged - stopping early!')
-                break
-
-            with open(learning_rate_tsv, 'a') as f:
-                f.write(f'{itr}\t{datetime.datetime.now():%H:%M:%S}\t{learning_rate}\t{loss_item}\n')
-
-        self.model.load_state_dict(model_state)
-        self.model.to(model_device)
-
-        log_line(log)
-        log.info(f'learning rate finder finished - plot {learning_rate_tsv}')
-        log_line(log)
-
-        return Path(learning_rate_tsv)
-
+         
     def train(self,
               base_path: Union[Path, str],
               evaluation_metric: EvaluationMetric = EvaluationMetric.MICRO_F1_SCORE,
@@ -114,7 +44,8 @@ def train(self,
               mini_batch_size: int = 32,
               max_epochs: int = 100,
               anneal_factor: float = 0.5,
-              patience: int = 4,
+              patience: int = 3,
+              anneal_against_train_loss: bool = True,
               train_with_dev: bool = False,
               monitor_train: bool = False,
               embeddings_in_memory: bool = True,
@@ -149,7 +80,7 @@ def train(self,
             optimizer.load_state_dict(self.optimizer_state)
 
         # annealing scheduler
-        anneal_mode = 'min' if train_with_dev else 'max'
+        anneal_mode = 'min' if anneal_against_train_loss else 'max'
         if isinstance(optimizer, (AdamW, SGDW)):
             scheduler = ReduceLRWDOnPlateau(optimizer, factor=anneal_factor,
                                             patience=patience, mode=anneal_mode,
@@ -167,8 +98,9 @@ def train(self,
         if train_with_dev:
             train_data.extend(self.corpus.dev)
 
-        loss_history = []
-        score_history = []
+        dev_score_history = []
+        dev_loss_history = []
+        train_loss_history = []
 
         # At any point you can hit Ctrl + C to break out of training early.
         try:
@@ -177,8 +109,7 @@ def train(self,
             for epoch in range(0 + self.epoch, max_epochs + self.epoch):
                 log_line(log)
 
-                # bad_epochs = scheduler.num_bad_epochs
-                bad_epochs = 0
+                bad_epochs = scheduler.num_bad_epochs
                 for group in optimizer.param_groups:
                     learning_rate = group['lr']
 
@@ -204,7 +135,7 @@ def train(self,
 
                 self.model.train()
 
-                current_loss: float = 0
+                train_loss: float = 0
                 seen_sentences = 0
                 modulo = max(1, int(len(batches) / 10))
 
@@ -217,23 +148,23 @@ def train(self,
                     optimizer.step()
 
                     seen_sentences += len(batch)
-                    current_loss += loss.item()
+                    train_loss += loss.item()
 
                     clear_embeddings(batch, also_clear_word_embeddings=not embeddings_in_memory)
 
                     if batch_no % modulo == 0:
                         log.info(f'epoch {epoch + 1} - iter {batch_no}/{len(batches)} - loss '
-                                 f'{current_loss / seen_sentences:.8f}')
+                                 f'{train_loss / seen_sentences:.8f}')
                         iteration = epoch * len(batches) + batch_no
                         if not param_selection_mode:
                             weight_extractor.extract_weights(self.model.state_dict(), iteration)
 
-                current_loss /= len(train_data)
+                train_loss /= len(train_data)
 
                 self.model.eval()
 
                 log_line(log)
-                log.info(f'EPOCH {epoch + 1}: lr {learning_rate:.4f} - bad epochs {bad_epochs}')
+                log.info(f'EPOCH {epoch + 1} done: loss {train_loss:.4f} - lr {learning_rate:.4f} - bad epochs {bad_epochs}')
 
                 dev_metric = None
                 dev_loss = '_'
@@ -247,9 +178,10 @@ def train(self,
                     dev_metric, dev_loss = self._calculate_evaluation_results_for(
                         'DEV', self.corpus.dev, evaluation_metric, embeddings_in_memory, mini_batch_size)
 
-                test_metric, test_loss = self._calculate_evaluation_results_for(
-                    'TEST', self.corpus.test, evaluation_metric, embeddings_in_memory, mini_batch_size,
-                    base_path / 'test.tsv')
+                if not param_selection_mode:
+                    test_metric, test_loss = self._calculate_evaluation_results_for(
+                        'TEST', self.corpus.test, evaluation_metric, embeddings_in_memory, mini_batch_size,
+                        base_path / 'test.tsv')
 
                 if not param_selection_mode:
                     with open(loss_txt, 'a') as f:
@@ -258,34 +190,36 @@ def train(self,
                         test_metric_str = test_metric.to_tsv() if test_metric is not None else Metric.to_empty_tsv()
                         f.write(
                             f'{epoch}\t{datetime.datetime.now():%H:%M:%S}\t{bad_epochs}\t{learning_rate:.4f}\t'
-                            f'{current_loss}\t{train_metric_str}\t{dev_loss}\t{dev_metric_str}\t_\t{test_metric_str}\n')
+                            f'{train_loss}\t{train_metric_str}\t{dev_loss}\t{dev_metric_str}\t_\t{test_metric_str}\n')
 
-                # anneal against train loss if training with dev, otherwise anneal against dev score
-                if train_with_dev:
-                    current_score = current_loss
-                else:
+                # calculate scores using dev data if available
+                dev_score = 0.
+                if not train_with_dev:
                     if evaluation_metric == EvaluationMetric.MACRO_ACCURACY:
-                        current_score = dev_metric.macro_avg_accuracy()
+                        dev_score = dev_metric.macro_avg_accuracy()
                     elif evaluation_metric == EvaluationMetric.MICRO_ACCURACY:
-                        current_score = dev_metric.micro_avg_accuracy()
+                        dev_score = dev_metric.micro_avg_accuracy()
                     elif evaluation_metric == EvaluationMetric.MACRO_F1_SCORE:
-                        current_score = dev_metric.macro_avg_f_score()
+                        dev_score = dev_metric.macro_avg_f_score()
                     else:
-                        current_score = dev_metric.micro_avg_f_score()
+                        dev_score = dev_metric.micro_avg_f_score()
+
+                    # append dev score to score history
+                    dev_score_history.append(dev_score)
+                    dev_loss_history.append(dev_loss.item())
+
+                # anneal against train loss if training with dev, otherwise anneal against dev score
+                current_score = train_loss if anneal_against_train_loss else dev_score
 
                 scheduler.step(current_score)
 
-                score_history.append(current_score)
-                if not train_with_dev:
-                    loss_history.append(dev_loss)
-                else:
-                    loss_history.append(current_loss)
+                train_loss_history.append(train_loss)
 
                 # if checkpoint is enable, save model at each epoch
                 if checkpoint and not param_selection_mode:
                     self.model.save_checkpoint(base_path / 'checkpoint.pt',
                                                optimizer.state_dict(), scheduler.state_dict(),
-                                               epoch + 1, current_loss)
+                                               epoch + 1, train_loss)
 
                 # if we use dev data, remember best model based on dev evaluation score
                 if not train_with_dev and not param_selection_mode and current_score == scheduler.best:
@@ -307,7 +241,10 @@ def train(self,
         if not param_selection_mode:
             final_score = self.final_test(base_path, embeddings_in_memory, evaluation_metric, mini_batch_size)
 
-        return {'test_score': final_score, 'score_history': score_history, 'loss_history': loss_history}
+        return {'test_score': final_score,
+                'dev_score_history': dev_score_history,
+                'train_loss_history': train_loss_history,
+                'dev_loss_history': dev_loss_history}
 
     def final_test(self,
                    base_path: Path,
@@ -502,3 +439,74 @@ def load_from_checkpoint(checkpoint_file: Path, model_type: str, corpus: Corpus,
                                 scheduler_state=checkpoint['scheduler_state_dict'])
 
         raise ValueError('Incorrect model type! Use one of the following: "SequenceTagger", "TextClassifier".')
+
+    def find_learning_rate(self,
+                           base_path: Union[Path, str],
+                           file_name: str = 'learning_rate.tsv',
+                           start_learning_rate: float = 1e-7,
+                           end_learning_rate: float = 10,
+                           iterations: int = 100,
+                           mini_batch_size: int = 32,
+                           stop_early: bool = True,
+                           smoothing_factor: float = 0.98,
+                           **kwargs
+                           ) -> Path:
+        best_loss = None
+        moving_avg_loss = 0
+
+        # cast string to Path
+        if type(base_path) is str:
+            base_path = Path(base_path)
+        learning_rate_tsv = init_output_file(base_path, file_name)
+
+        with open(learning_rate_tsv, 'a') as f:
+            f.write('ITERATION\tTIMESTAMP\tLEARNING_RATE\tTRAIN_LOSS\n')
+
+        optimizer = self.optimizer(self.model.parameters(), lr=start_learning_rate, **kwargs)
+
+        train_data = self.corpus.train
+        random.shuffle(train_data)
+        batches = [train_data[x:x + mini_batch_size] for x in range(0, len(train_data), mini_batch_size)][:iterations]
+
+        scheduler = ExpAnnealLR(optimizer, end_learning_rate, iterations)
+
+        model_state = self.model.state_dict()
+        model_device = next(self.model.parameters()).device
+        self.model.train()
+
+        for itr, batch in enumerate(batches):
+            loss = self.model.forward_loss(batch)
+
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5.0)
+            optimizer.step()
+            scheduler.step()
+            learning_rate = scheduler.get_lr()[0]
+
+            loss_item = loss.item()
+            if itr == 0:
+                best_loss = loss_item
+            else:
+                if smoothing_factor > 0:
+                    moving_avg_loss = smoothing_factor * moving_avg_loss + (1 - smoothing_factor) * loss_item
+                    loss_item = moving_avg_loss / (1 - smoothing_factor ** (itr+1))
+                if loss_item < best_loss:
+                    best_loss = loss
+
+            if stop_early and (loss_item > 4 * best_loss or torch.isnan(loss)):
+                log_line(log)
+                log.info('loss diverged - stopping early!')
+                break
+
+            with open(learning_rate_tsv, 'a') as f:
+                f.write(f'{itr}\t{datetime.datetime.now():%H:%M:%S}\t{learning_rate}\t{loss_item}\n')
+
+        self.model.load_state_dict(model_state)
+        self.model.to(model_device)
+
+        log_line(log)
+        log.info(f'learning rate finder finished - plot {learning_rate_tsv}')
+        log_line(log)
+
+        return Path(learning_rate_tsv)
\ No newline at end of file


Test patch: