Repository: jrzaurin/pytorch-widedeep (language: python)

Issue:
EarlyStopping does not store and restore the model
Hi. I found a confusing issue in `EarlyStopping` and I think it is an essential bug. When using `EarlyStopping` without an additional `ModelCheckpoint`, the best weight will not be restored. The bug is made up of three mistakes.

Firstly, the `state_dict()` object changes while the model is training because it is not a **deep** copy (even a shallow copy `state_dict().copy()` will not work). See [this discussion](https://discuss.pytorch.org/t/state-dict-changes-with-the-model/47547/5). So the following lines are actually not storing anything.

https://github.com/jrzaurin/pytorch-widedeep/blob/9acfceca015a7553f181fb86f99a0d8ed7b5e798/pytorch_widedeep/callbacks.py#L659-L660

Secondly, if `n_epoch` is reached without triggering early stopping, the best weight is not restored, because nothing is done in `on_epoch_end` or

https://github.com/jrzaurin/pytorch-widedeep/blob/9acfceca015a7553f181fb86f99a0d8ed7b5e798/pytorch_widedeep/callbacks.py#L671-L673

Thirdly, if no `ModelCheckpoint` is used, after training ends,

https://github.com/jrzaurin/pytorch-widedeep/blob/9acfceca015a7553f181fb86f99a0d8ed7b5e798/pytorch_widedeep/training/trainer.py#L532

The method `_restore_best_weights` mistakenly thinks that the best weight is restored by `EarlyStopping` (no matter whether early stopping is triggered because of the above two points)

https://github.com/jrzaurin/pytorch-widedeep/blob/9acfceca015a7553f181fb86f99a0d8ed7b5e798/pytorch_widedeep/training/_base_trainer.py#L133-L144

The above issue would not happen if `ModelCheckpoint` is used because the best weight is correctly stored and restored.

Here is my current solution that only changes two lines, which can be your reference.

```python
from pytorch_widedeep.callbacks import EarlyStopping as ES
from copy import deepcopy as cp

class EarlyStopping(ES):
    def on_epoch_end(
        self, epoch: int, logs: Optional[Dict] = None, metric: Optional[float] = None
    ):
        current = self.get_monitor_value(logs)
        if current is None:
            return

        if self.monitor_op(current - self.min_delta, self.best):
            self.best = current
            self.wait = 0
            if self.restore_best_weights:
                self.state_dict = cp(self.model.state_dict()) # here
        else:
            self.wait += 1
            if self.wait >= self.patience:
                self.stopped_epoch = epoch
                self.trainer.early_stop = True
                if self.restore_best_weights:
                    if self.verbose > 0:
                        print("Restoring model weights from the end of the best epoch")
                    self.model.load_state_dict(self.state_dict)

    def on_train_end(self, logs: Optional[Dict] = None):
        super(EarlyStopping, self).on_train_end(logs)
        self.model.load_state_dict(self.state_dict) # and here
```

EarlyStopping does not store and restore the model
Hi. I found a confusing issue in `EarlyStopping` and I think it is an essential bug. When using `EarlyStopping` without an additional `ModelCheckpoint`, the best weight will not be restored. The bug is made up of three mistakes.

Firstly, the `state_dict()` object changes while the model is training because it is not a **deep** copy (even a shallow copy `state_dict().copy()` will not work). See [this discussion](https://discuss.pytorch.org/t/state-dict-changes-with-the-model/47547/5). So the following lines are actually not storing anything.

https://github.com/jrzaurin/pytorch-widedeep/blob/9acfceca015a7553f181fb86f99a0d8ed7b5e798/pytorch_widedeep/callbacks.py#L659-L660

Secondly, if `n_epoch` is reached without triggering early stopping, the best weight is not restored, because nothing is done in `on_epoch_end` or

https://github.com/jrzaurin/pytorch-widedeep/blob/9acfceca015a7553f181fb86f99a0d8ed7b5e798/pytorch_widedeep/callbacks.py#L671-L673

Thirdly, if no `ModelCheckpoint` is used, after training ends,

https://github.com/jrzaurin/pytorch-widedeep/blob/9acfceca015a7553f181fb86f99a0d8ed7b5e798/pytorch_widedeep/training/trainer.py#L532

The method `_restore_best_weights` mistakenly thinks that the best weight is restored by `EarlyStopping` (no matter whether early stopping is triggered because of the above two points)

https://github.com/jrzaurin/pytorch-widedeep/blob/9acfceca015a7553f181fb86f99a0d8ed7b5e798/pytorch_widedeep/training/_base_trainer.py#L133-L144

The above issue would not happen if `ModelCheckpoint` is used because the best weight is correctly stored and restored.

Here is my current solution that only changes two lines, which can be your reference.

```python
from pytorch_widedeep.callbacks import EarlyStopping as ES
from copy import deepcopy as cp

class EarlyStopping(ES):
    def on_epoch_end(
        self, epoch: int, logs: Optional[Dict] = None, metric: Optional[float] = None
    ):
        current = self.get_monitor_value(logs)
        if current is None:
            return

        if self.monitor_op(current - self.min_delta, self.best):
            self.best = current
            self.wait = 0
            if self.restore_best_weights:
                self.state_dict = cp(self.model.state_dict()) # here
        else:
            self.wait += 1
            if self.wait >= self.patience:
                self.stopped_epoch = epoch
                self.trainer.early_stop = True
                if self.restore_best_weights:
                    if self.verbose > 0:
                        print("Restoring model weights from the end of the best epoch")
                    self.model.load_state_dict(self.state_dict)

    def on_train_end(self, logs: Optional[Dict] = None):
        super(EarlyStopping, self).on_train_end(logs)
        self.model.load_state_dict(self.state_dict) # and here
```



Discussion / hints:
hi @LuoXueling 

I will look to it asap. 

In any case I am wondering, would you want to contribute and fix it? 🙂

Otherwise, thanks for the issue and the detail explanation ;) 
Sure. I will make a PR soon.
hi @LuoXueling 

I will look to it asap. 

In any case I am wondering, would you want to contribute and fix it? 🙂

Otherwise, thanks for the issue and the detail explanation ;) 
Sure. I will make a PR soon.

Fix patch:
diff --git a/pytorch_widedeep/callbacks.py b/pytorch_widedeep/callbacks.py
index 4202086a..b721bac4 100644
--- a/pytorch_widedeep/callbacks.py
+++ b/pytorch_widedeep/callbacks.py
@@ -6,6 +6,7 @@
 import os
 import datetime
 import warnings
+import copy
 
 import numpy as np
 import torch
@@ -657,20 +658,20 @@ def on_epoch_end(
             self.best = current
             self.wait = 0
             if self.restore_best_weights:
-                self.state_dict = self.model.state_dict()
+                self.state_dict = copy.deepcopy(self.model.state_dict())
         else:
             self.wait += 1
             if self.wait >= self.patience:
                 self.stopped_epoch = epoch
                 self.trainer.early_stop = True
-                if self.restore_best_weights:
-                    if self.verbose > 0:
-                        print("Restoring model weights from the end of the best epoch")
-                    self.model.load_state_dict(self.state_dict)
 
     def on_train_end(self, logs: Optional[Dict] = None):
         if self.stopped_epoch > 0 and self.verbose > 0:
             print("Epoch %05d: early stopping" % (self.stopped_epoch + 1))
+        if self.restore_best_weights and self.state_dict is not None:
+            if self.verbose > 0:
+                print("Restoring model weights from the end of the best epoch")
+            self.model.load_state_dict(self.state_dict)
 
     def get_monitor_value(self, logs):
         monitor_value = logs.get(self.monitor)


Test patch:
diff --git a/tests/test_model_functioning/test_callbacks.py b/tests/test_model_functioning/test_callbacks.py
index ee47dcdc..ac7fa08b 100644
--- a/tests/test_model_functioning/test_callbacks.py
+++ b/tests/test_model_functioning/test_callbacks.py
@@ -479,3 +479,30 @@ def test_early_stopping_get_state():
     shutil.rmtree("tests/test_model_functioning/early_stopping/")
 
     assert no_trainer and no_model
+
+
+def test_early_stopping_restore_state():
+    # min_delta is large, so the early stopping condition will never be met except for the first epoch.
+    early_stopping = EarlyStopping(
+        restore_best_weights=True, min_delta=1000, patience=1000
+    )
+    trainer_tt = Trainer(
+        model,
+        objective="regression",
+        callbacks=[early_stopping],
+        verbose=0,
+    )
+    trainer_tt.fit(
+        X_train={"X_wide": X_wide, "X_tab": X_tab, "target": target},
+        X_val={"X_wide": X_wide_val, "X_tab": X_tab_val, "target": target_val},
+        target=target,
+        n_epochs=2,
+        batch_size=16,
+    )
+    assert early_stopping.wait > 0
+    # so early stopping is not triggered, but is over-fitting.
+    pred_val = trainer_tt.predict(X_test={"X_wide": X_wide_val, "X_tab": X_tab_val})
+    restored_metric = trainer_tt.loss_fn(
+        torch.tensor(pred_val), torch.tensor(target_val)
+    ).item()
+    assert np.allclose(restored_metric, early_stopping.best)