Repository: flairNLP/flair

Issue:
When the dev is with a perfect score, consider a smaller loss an improvment
Hello,
I'm training a very simple dataset, and realized when it archive a perfect score, it do not consider the loss improvement as a good epoch. 

```
2020-02-03 04:45:49,242 EPOCH 42 done: loss 3.9541 - lr 0.0250
2020-02-03 04:46:11,877 DEV : loss 2.556217670440674 - score 1.0
2020-02-03 04:46:12,227 BAD EPOCHS (no improvement): 1
2020-02-03 04:46:18,934 ----------------------------------------------------------------------------------------------------
2020-02-03 04:46:20,809 epoch 43 - iter 1/15 - loss 0.11718750 - samples/sec: 0.55
2020-02-03 04:46:25,734 epoch 43 - iter 2/15 - loss 0.19140625 - samples/sec: 0.20
2020-02-03 04:46:26,519 epoch 43 - iter 3/15 - loss 0.13997396 - samples/sec: 1.32
2020-02-03 04:46:31,470 epoch 43 - iter 4/15 - loss 0.11376953 - samples/sec: 0.20
2020-02-03 04:46:35,815 epoch 43 - iter 5/15 - loss 0.12031250 - samples/sec: 0.23
2020-02-03 04:46:38,706 epoch 43 - iter 6/15 - loss 4.69742839 - samples/sec: 0.35
2020-02-03 04:46:41,963 epoch 43 - iter 7/15 - loss 4.03069196 - samples/sec: 0.31
2020-02-03 04:46:46,578 epoch 43 - iter 8/15 - loss 3.53417969 - samples/sec: 0.22
2020-02-03 04:46:48,116 epoch 43 - iter 9/15 - loss 3.14615885 - samples/sec: 0.66
2020-02-03 04:46:49,544 epoch 43 - iter 10/15 - loss 2.84052734 - samples/sec: 0.71
2020-02-03 04:46:51,052 epoch 43 - iter 11/15 - loss 2.58540483 - samples/sec: 0.67
2020-02-03 04:47:12,472 epoch 43 - iter 12/15 - loss 3.07373047 - samples/sec: 0.05
2020-02-03 04:47:33,290 epoch 43 - iter 13/15 - loss 2.84990986 - samples/sec: 0.05
2020-02-03 04:47:35,578 epoch 43 - iter 14/15 - loss 2.64885603 - samples/sec: 0.44
2020-02-03 04:47:37,917 epoch 43 - iter 15/15 - loss 2.50240885 - samples/sec: 0.43
2020-02-03 04:47:37,938 ----------------------------------------------------------------------------------------------------
2020-02-03 04:47:37,939 EPOCH 43 done: loss 2.5024 - lr 0.0250
2020-02-03 04:47:59,749 DEV : loss 1.2600260972976685 - score 1.0
2020-02-03 04:48:00,162 BAD EPOCHS (no improvement): 2
```


Discussion:
Recently. I had similar "problem". The `best-model.pt` and `test.tsv` are overwritten if the score in the next epoch is the same as the best. So, instead of having a model after 16 epochs I have a model after the last epoch. The above mentioned strategy could solve my problem.
```
EPOCH	TIMESTAMP	BAD_EPOCHS	LEARNING_RATE	TRAIN_LOSS	DEV_LOSS	DEV_PRECISION	DEV_RECALL	DEV_F1	TEST_LOSS	TEST_PRECISION	TEST_RECALL	TEST_F1
1	21:09:11	0	0.1000	4.870719216444312	0.3633814752101898	0.9989	0.9989	0.9989	0.6925714612007141	0.9984	0.9984	0.9984
2	21:16:12	0	0.1000	0.7478587668721038	0.23151203989982605	0.9994	0.9994	0.9994	0.41518688201904297	0.999	0.999	0.999
3	21:23:14	1	0.1000	0.5148863803238759	0.21521516144275665	0.9994	0.9994	0.9994	0.31954842805862427	0.9994	0.9994	0.9994
4	21:30:20	2	0.1000	0.4522323644790712	0.23185299336910248	0.9994	0.9994	0.9994	0.382921040058136	0.9991	0.9991	0.9991
5	21:37:28	3	0.1000	0.4250520358581354	0.20026066899299622	0.9994	0.9994	0.9994	0.31172817945480347	0.9993	0.9993	0.9993
6	21:44:27	0	0.1000	0.36263318463127214	0.2035578191280365	0.9995	0.9995	0.9995	0.31626230478286743	0.9993	0.9993	0.9993
7	21:51:35	1	0.1000	0.3679700958551747	0.18492664396762848	0.9995	0.9995	0.9995	0.2523406147956848	0.9994	0.9994	0.9994
8	21:58:44	2	0.1000	0.3587756136266312	0.1910945475101471	0.9995	0.9995	0.9995	0.2961016893386841	0.9994	0.9994	0.9994
9	22:05:54	3	0.1000	0.31385309890945357	0.17560453712940216	0.9995	0.9995	0.9995	0.2709144353866577	0.9994	0.9994	0.9994
10	22:13:01	4	0.1000	0.3031430102872475	0.17703454196453094	0.9995	0.9995	0.9995	0.2912594676017761	0.9994	0.9994	0.9994
11	22:20:02	5	0.1000	0.2935463251060385	0.1718766689300537	0.9995	0.9995	0.9995	0.2600330710411072	0.9994	0.9994	0.9994
12	22:27:10	6	0.1000	0.2980318038475395	0.16834411025047302	0.9995	0.9995	0.9995	0.24408173561096191	0.9994	0.9994	0.9994
13	22:34:18	1	0.0500	0.26744764954736916	0.17041155695915222	0.9995	0.9995	0.9995	0.24594569206237793	0.9994	0.9994	0.9994
14	22:41:27	2	0.0500	0.2472817079462234	0.16346509754657745	0.9995	0.9995	0.9995	0.25140634179115295	0.9994	0.9994	0.9994
15	22:48:37	3	0.0500	0.24876306218282618	0.16812703013420105	0.9995	0.9995	0.9995	0.23818974196910858	0.9994	0.9994	0.9994
16	22:55:37	0	0.0500	0.2389219280320032	0.160366028547287	0.9996	0.9996	0.9996	0.2426607608795166	0.9994	0.9994	0.9994
17	23:02:48	1	0.0500	0.24479633510702908	0.1646646112203598	0.9996	0.9996	0.9996	0.26128894090652466	0.9994	0.9994	0.9994
18	23:09:59	2	0.0500	0.327062593139831	0.17466337978839874	0.9995	0.9995	0.9995	0.25381121039390564	0.9995	0.9995	0.9995
19	23:17:08	3	0.0500	0.21864327849826404	0.16761788725852966	0.9996	0.9996	0.9996	0.25565916299819946	0.9994	0.9994	0.9994
20	23:24:18	4	0.0500	0.2646045185158355	0.17045339941978455	0.9996	0.9996	0.9996	0.26664477586746216	0.9993	0.9993	0.9993
21	23:31:19	5	0.0500	0.22782475110327843	0.1694835126399994	0.9996	0.9996	0.9996	0.24179255962371826	0.9994	0.9994	0.9994
22	23:38:29	6	0.0500	0.20819690105545247	0.1753673255443573	0.9996	0.9996	0.9996	0.27125662565231323	0.9993	0.9993	0.9993
23	23:45:40	1	0.0250	0.22172029611497823	0.164693683385849	0.9996	0.9996	0.9996	0.2551299035549164	0.9993	0.9993	0.9993
24	23:52:49	2	0.0250	0.199604090346214	0.17024561762809753	0.9996	0.9996	0.9996	0.26467064023017883	0.9993	0.9993	0.9993
25	00:00:00	3	0.0250	0.19735614704613638	0.16565625369548798	0.9996	0.9996	0.9996	0.24156562983989716	0.9994	0.9994	0.9994
26	00:07:00	4	0.0250	0.20577437083910008	0.18168862164020538	0.9996	0.9996	0.9996	0.2920682728290558	0.9994	0.9994	0.9994
27	00:14:10	5	0.0250	0.20145670661438417	0.1726815402507782	0.9996	0.9996	0.9996	0.2629989981651306	0.9993	0.9993	0.9993
28	00:21:20	6	0.0250	0.19017612250527927	0.17380431294441223	0.9995	0.9995	0.9995	0.24805815517902374	0.9995	0.9995	0.9995
29	00:28:30	1	0.0125	0.18600339850675154	0.1714184582233429	0.9995	0.9995	0.9995	0.24996189773082733	0.9994	0.9994	0.9994
30	00:35:39	2	0.0125	0.18729554356472328	0.17530648410320282	0.9996	0.9996	0.9996	0.2646336257457733	0.9993	0.9993	0.9993
31	00:42:39	3	0.0125	0.18741807277643246	0.1716986447572708	0.9995	0.9995	0.9995	0.26211386919021606	0.9993	0.9993	0.9993
32	00:49:48	4	0.0125	0.18372498586982883	0.17321093380451202	0.9995	0.9995	0.9995	0.25690072774887085	0.9994	0.9994	0.9994
33	00:56:58	5	0.0125	0.1824594216771645	0.17323888838291168	0.9996	0.9996	0.9996	0.256899893283844	0.9994	0.9994	0.9994
34	01:04:07	6	0.0125	0.1891614758707706	0.17295822501182556	0.9996	0.9996	0.9996	0.2543550133705139	0.9994	0.9994	0.9994
35	01:11:16	1	0.0063	0.17181612054506937	0.1717613935470581	0.9996	0.9996	0.9996	0.25098150968551636	0.9994	0.9994	0.9994
36	01:18:16	2	0.0063	0.18608331336046602	0.17340107262134552	0.9996	0.9996	0.9996	0.2596213221549988	0.9993	0.9993	0.9993
37	01:25:28	3	0.0063	0.16709851380130245	0.17306150496006012	0.9996	0.9996	0.9996	0.24840085208415985	0.9995	0.9995	0.9995
38	01:32:38	4	0.0063	0.17404617671624267	0.17438027262687683	0.9996	0.9996	0.9996	0.255127489566803	0.9993	0.9993	0.9993
39	01:39:48	5	0.0063	0.18388921078970724	0.17367535829544067	0.9996	0.9996	0.9996	0.2543256878852844	0.9994	0.9994	0.9994
40	01:46:58	6	0.0063	0.2526034320590913	0.17365668714046478	0.9996	0.9996	0.9996	0.25782349705696106	0.9993	0.9993	0.9993
41	01:53:58	1	0.0031	0.17762035531646544	0.17282190918922424	0.9996	0.9996	0.9996	0.2573838233947754	0.9993	0.9993	0.9993
42	02:01:07	2	0.0031	0.1892923788404209	0.17292071878910065	0.9996	0.9996	0.9996	0.25628048181533813	0.9993	0.9993	0.9993
43	02:08:15	3	0.0031	0.17523770715438886	0.17280033230781555	0.9996	0.9996	0.9996	0.25576916337013245	0.9994	0.9994	0.9994
44	02:15:22	4	0.0031	0.21022769690740226	0.17306119203567505	0.9996	0.9996	0.9996	0.2550627589225769	0.9994	0.9994	0.9994
45	02:22:22	5	0.0031	0.17243955327388477	0.17395059764385223	0.9996	0.9996	0.9996	0.2589952349662781	0.9994	0.9994	0.9994
46	02:29:05	6	0.0031	0.17081469395766558	0.17403607070446014	0.9996	0.9996	0.9996	0.25671932101249695	0.9994	0.9994	0.9994
47	02:35:56	1	0.0016	0.16745367223280097	0.1738416999578476	0.9996	0.9996	0.9996	0.2579036355018616	0.9994	0.9994	0.9994
48	02:42:48	2	0.0016	0.16530116485192045	0.1746242493391037	0.9996	0.9996	0.9996	0.2589937150478363	0.9994	0.9994	0.9994
49	02:49:40	3	0.0016	0.1747354833599758	0.17429324984550476	0.9996	0.9996	0.9996	0.2571435272693634	0.9994	0.9994	0.9994
50	02:56:31	4	0.0016	0.1711603099166757	0.17396044731140137	0.9996	0.9996	0.9996	0.2544676959514618	0.9994	0.9994	0.9994
```
Good point. Maybe we should change the default annealing strategy to always go against dev loss instead of dev score? This would solve this issue and score and loss are correlated anyway.
No, the loss does not always correspond to metric, e.g. F-beta.
Yes you're right - in that case we'd need a special scheduler that uses default score normally and uses loss only as a fallback if the same best score appears in multiple epochs.

Resolving patch:
diff --git a/flair/__init__.py b/flair/__init__.py
--- a/flair/__init__.py
+++ b/flair/__init__.py
@@ -20,6 +20,7 @@
 from . import visual
 from . import trainers
 from . import nn
+from .training_utils import AnnealOnPlateau
 
 import logging.config
 
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 @@
+import copy
 import logging
 from pathlib import Path
 from typing import List, Union
@@ -7,7 +8,6 @@
 import inspect
 
 import torch
-from torch.optim.lr_scheduler import ReduceLROnPlateau
 from torch.optim.sgd import SGD
 from torch.utils.data.dataset import ConcatDataset
 
@@ -28,6 +28,7 @@
     add_file_handler,
     Result,
     store_embeddings,
+    AnnealOnPlateau,
 )
 from flair.models import SequenceTagger
 import random
@@ -65,6 +66,7 @@ def train(
         mini_batch_size: int = 32,
         mini_batch_chunk_size: int = None,
         max_epochs: int = 100,
+        scheduler = AnnealOnPlateau,
         anneal_factor: float = 0.5,
         patience: int = 3,
         min_learning_rate: float = 0.0001,
@@ -220,7 +222,7 @@ def train(
         # minimize training loss if training with dev data, else maximize dev score
         anneal_mode = "min" if train_with_dev else "max"
 
-        scheduler: ReduceLROnPlateau = ReduceLROnPlateau(
+        lr_scheduler = scheduler(
             optimizer,
             factor=anneal_factor,
             patience=patience,
@@ -256,8 +258,8 @@ def train(
             for self.epoch in range(self.epoch + 1, max_epochs + 1):
                 log_line(log)
 
-                last_epoch_model_state_dict = self.model.state_dict()
-                best_state_dict = self.model.state_dict()
+                if anneal_with_prestarts:
+                    last_epoch_model_state_dict = copy.deepcopy(self.model.state_dict())
 
                 if eval_on_train_shuffle:
                     train_part_indices = list(range(self.corpus.train))
@@ -276,20 +278,21 @@ def train(
 
                 # reload last best model if annealing with restarts is enabled
                 if (
-                    learning_rate != previous_learning_rate
+                    (anneal_with_restarts or anneal_with_prestarts)
+                    and learning_rate != previous_learning_rate
                     and (base_path / "best-model.pt").exists()
                 ):
-                    log.info("resetting to best model")
                     if anneal_with_restarts:
+                        log.info("resetting to best model")
                         self.model.load_state_dict(
                             self.model.load(base_path / "best-model.pt").state_dict()
                         )
                     if anneal_with_prestarts:
+                        log.info("resetting to pre-best model")
                         self.model.load_state_dict(
                             self.model.load(base_path / "pre-best-model.pt").state_dict()
                         )
 
-
                 previous_learning_rate = learning_rate
 
                 # stop training if learning rate becomes too small
@@ -470,14 +473,17 @@ def train(
                             "test_score", test_eval_result.main_score, self.epoch
                         )
 
-                # determine learning rate annealing through scheduler
-                scheduler.step(current_score)
+                # determine learning rate annealing through scheduler. Use auxiliary metric for AnnealOnPlateau
+                if not train_with_dev and isinstance(lr_scheduler, AnnealOnPlateau):
+                    lr_scheduler.step(current_score, dev_loss)
+                else:
+                    lr_scheduler.step(current_score)
 
                 train_loss_history.append(train_loss)
 
                 # determine bad epoch number
                 try:
-                    bad_epochs = scheduler.num_bad_epochs
+                    bad_epochs = lr_scheduler.num_bad_epochs
                 except:
                     bad_epochs = 0
                 for group in optimizer.param_groups:
@@ -537,15 +543,17 @@ def train(
                 if (
                     (not train_with_dev or anneal_with_restarts or anneal_with_prestarts)
                     and not param_selection_mode
-                    and current_score == scheduler.best
+                    and current_score == lr_scheduler.best
+                    and bad_epochs == 0
                 ):
                     print("saving best model")
                     self.model.save(base_path / "best-model.pt")
-                    current_state_dict = self.model.state_dict()
-                    self.model.load_state_dict(last_epoch_model_state_dict)
+
                     if anneal_with_prestarts:
+                        current_state_dict = self.model.state_dict()
+                        self.model.load_state_dict(last_epoch_model_state_dict)
                         self.model.save(base_path / "pre-best-model.pt")
-                    self.model.load_state_dict(current_state_dict)
+                        self.model.load_state_dict(current_state_dict)
 
             # if we do not use dev data for model selection, save final model
             if save_final_model and not param_selection_mode:
diff --git a/flair/training_utils.py b/flair/training_utils.py
--- a/flair/training_utils.py
+++ b/flair/training_utils.py
@@ -3,9 +3,12 @@
 import logging
 from collections import defaultdict
 from enum import Enum
+from math import inf
 from pathlib import Path
 from typing import List
 
+from torch.optim import Optimizer
+
 import flair
 from flair.data import Dictionary, Sentence
 from functools import reduce
@@ -294,6 +297,172 @@ def _init_weights_index(self, key, state_dict, weights_to_watch):
         self.weights_dict[key] = indices
 
 
+class AnnealOnPlateau(object):
+    """This class is a modification of
+    torch.optim.lr_scheduler.ReduceLROnPlateau that enables
+    setting an "auxiliary metric" to break ties.
+
+    Reduce learning rate when a metric has stopped improving.
+    Models often benefit from reducing the learning rate by a factor
+    of 2-10 once learning stagnates. This scheduler reads a metrics
+    quantity and if no improvement is seen for a 'patience' number
+    of epochs, the learning rate is reduced.
+
+    Args:
+        optimizer (Optimizer): Wrapped optimizer.
+        mode (str): One of `min`, `max`. In `min` mode, lr will
+            be reduced when the quantity monitored has stopped
+            decreasing; in `max` mode it will be reduced when the
+            quantity monitored has stopped increasing. Default: 'min'.
+        factor (float): Factor by which the learning rate will be
+            reduced. new_lr = lr * factor. Default: 0.1.
+        patience (int): Number of epochs with no improvement after
+            which learning rate will be reduced. For example, if
+            `patience = 2`, then we will ignore the first 2 epochs
+            with no improvement, and will only decrease the LR after the
+            3rd epoch if the loss still hasn't improved then.
+            Default: 10.
+        verbose (bool): If ``True``, prints a message to stdout for
+            each update. Default: ``False``.
+        cooldown (int): Number of epochs to wait before resuming
+            normal operation after lr has been reduced. Default: 0.
+        min_lr (float or list): A scalar or a list of scalars. A
+            lower bound on the learning rate of all param groups
+            or each group respectively. Default: 0.
+        eps (float): Minimal decay applied to lr. If the difference
+            between new and old lr is smaller than eps, the update is
+            ignored. Default: 1e-8.
+
+    Example:
+        >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
+        >>> scheduler = ReduceLROnPlateau(optimizer, 'min')
+        >>> for epoch in range(10):
+        >>>     train(...)
+        >>>     val_loss = validate(...)
+        >>>     # Note that step should be called after validate()
+        >>>     scheduler.step(val_loss)
+    """
+
+    def __init__(self, optimizer, mode='min', aux_mode='min', factor=0.1, patience=10,
+                 verbose=False, cooldown=0, min_lr=0, eps=1e-8):
+
+        if factor >= 1.0:
+            raise ValueError('Factor should be < 1.0.')
+        self.factor = factor
+
+        # Attach optimizer
+        if not isinstance(optimizer, Optimizer):
+            raise TypeError('{} is not an Optimizer'.format(
+                type(optimizer).__name__))
+        self.optimizer = optimizer
+
+        if isinstance(min_lr, list) or isinstance(min_lr, tuple):
+            if len(min_lr) != len(optimizer.param_groups):
+                raise ValueError("expected {} min_lrs, got {}".format(
+                    len(optimizer.param_groups), len(min_lr)))
+            self.min_lrs = list(min_lr)
+        else:
+            self.min_lrs = [min_lr] * len(optimizer.param_groups)
+
+        self.patience = patience
+        self.verbose = verbose
+        self.cooldown = cooldown
+        self.cooldown_counter = 0
+        self.mode = mode
+        self.aux_mode = aux_mode
+        self.best = None
+        self.best_aux = None
+        self.num_bad_epochs = None
+        self.mode_worse = None  # the worse value for the chosen mode
+        self.eps = eps
+        self.last_epoch = 0
+        self._init_is_better(mode=mode)
+        self._reset()
+
+    def _reset(self):
+        """Resets num_bad_epochs counter and cooldown counter."""
+        self.best = self.mode_worse
+        self.cooldown_counter = 0
+        self.num_bad_epochs = 0
+
+    def step(self, metric, auxiliary_metric = None):
+        # convert `metrics` to float, in case it's a zero-dim Tensor
+        current = float(metric)
+        epoch = self.last_epoch + 1
+        self.last_epoch = epoch
+
+        is_better = False
+
+        if self.mode == 'min':
+            if current < self.best:
+                is_better = True
+
+        if self.mode == 'max':
+            if current > self.best:
+                is_better = True
+
+        if current == self.best and auxiliary_metric:
+            current_aux = float(auxiliary_metric)
+            if self.aux_mode == 'min':
+                if current_aux < self.best_aux:
+                    is_better = True
+
+            if self.aux_mode == 'max':
+                if current_aux > self.best_aux:
+                    is_better = True
+
+        if is_better:
+            self.best = current
+            if auxiliary_metric:
+                self.best_aux = auxiliary_metric
+            self.num_bad_epochs = 0
+        else:
+            self.num_bad_epochs += 1
+
+        if self.in_cooldown:
+            self.cooldown_counter -= 1
+            self.num_bad_epochs = 0  # ignore any bad epochs in cooldown
+
+        if self.num_bad_epochs > self.patience:
+            self._reduce_lr(epoch)
+            self.cooldown_counter = self.cooldown
+            self.num_bad_epochs = 0
+
+        self._last_lr = [group['lr'] for group in self.optimizer.param_groups]
+
+    def _reduce_lr(self, epoch):
+        for i, param_group in enumerate(self.optimizer.param_groups):
+            old_lr = float(param_group['lr'])
+            new_lr = max(old_lr * self.factor, self.min_lrs[i])
+            if old_lr - new_lr > self.eps:
+                param_group['lr'] = new_lr
+                if self.verbose:
+                    print('Epoch {:5d}: reducing learning rate'
+                          ' of group {} to {:.4e}.'.format(epoch, i, new_lr))
+
+    @property
+    def in_cooldown(self):
+        return self.cooldown_counter > 0
+
+    def _init_is_better(self, mode):
+        if mode not in {'min', 'max'}:
+            raise ValueError('mode ' + mode + ' is unknown!')
+
+        if mode == 'min':
+            self.mode_worse = inf
+        else:  # mode == 'max':
+            self.mode_worse = -inf
+
+        self.mode = mode
+
+    def state_dict(self):
+        return {key: value for key, value in self.__dict__.items() if key != 'optimizer'}
+
+    def load_state_dict(self, state_dict):
+        self.__dict__.update(state_dict)
+        self._init_is_better(mode=self.mode, threshold=self.threshold, threshold_mode=self.threshold_mode)
+
+
 def init_output_file(base_path: Path, file_name: str) -> Path:
     """
     Creates a local file.


Test patch: