Repository: Lightning-AI/pytorch-lightning

Issue:
Checkpoints not working
## 🐛 Bug

Default `checkpoint_callback` in `Trainer()` does not work so model's checkpoints are not saved. 

### To Reproduce

Steps to reproduce the behavior:

I first created a simple implementation of a `LightningModule`. This contains:
```
class Model(pl.LightningModule):
   ...
   def validation_step(self, batch, batch_idx):
      x, y = batch
      y_hat = self(x)
      loss = F.cross_entropy(y_hat, y)
      result = pl.EvalResult()
      result.log('val_loss', loss, prog_bar=True, on_step=False, on_epoch=True)
      return result
   ...
```
when I run it with 
```
m = Model()
trainer.fit(model) # using default settings
```
I get 
```
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/utilities/distributed.py:37: RuntimeWarning: The metric you returned None must be a `torch.Tensor` instance, checkpoint not saved HINT: what is the value of loss in validation_epoch_end()?
  warnings.warn(*args, **kwargs)
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/utilities/distributed.py:37: RuntimeWarning: Can save best model only with loss available, skipping.
  warnings.warn(*args, **kwargs)
```

Following the hint, I then implemented the `validation_epoch_end()` method. I was surprised I had to, given that most examples in the doc work without it and I assumed it was used when additional control on the validation procedure was needed. Perhaps, I would disable all callbacks by defaults to avoid the need to implement additional methods at first. Anyways, now my model looks like:
```
class Model(pl.LightningModule):
   ...
   def validation_step(self, batch, batch_idx):
      x, y = batch
      y_hat = self(x)
      loss = F.cross_entropy(y_hat, y)
      result = pl.EvalResult()
      result.log('val_loss', loss, prog_bar=True, on_step=False, on_epoch=True)
      return result

   def validation_epoch_end(self, outputs):
      avg_loss = torch.stack([x for x in outputs['val_loss']]).mean()
      return {"val_loss": avg_loss}
   ...
```
In the above, I followed https://github.com/PyTorchLightning/pytorch-lightning/issues/1153#issuecomment-599792149 but got:
```
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/evaluate_loop.py in log_epoch_metrics(self, eval_results)
    148             if isinstance(eval_results, list):
    149                 for eval_result in eval_results:
--> 150                     self.trainer.callback_metrics = eval_result.callback_metrics
    151             else:
    152                 self.trainer.callback_metrics = eval_results.callback_metrics

AttributeError: 'dict' object has no attribute 'callback_metrics'
```
Thus, I tried to use the new `Result` container:
```
class Model(pl.LightningModule):
   ...
   def validation_step(self, batch, batch_idx):
      x, y = batch
      y_hat = self(x)
      loss = F.cross_entropy(y_hat, y)
      result = pl.EvalResult()
      result.log('val_loss', loss, prog_bar=True, on_step=False, on_epoch=True)
      return result

   def validation_epoch_end(self, outputs):
      loss = torch.stack([x for x in outputs['val_loss']]).mean()
      result = pl.EvalResult()
      result.log('val_loss', loss)
      return result
   ...
```

When I fit the model again, I get:
```
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/utilities/distributed.py:37: RuntimeWarning: The metric you returned None must be a `torch.Tensor` instance, checkpoint not saved HINT: what is the value of loss in validation_epoch_end()?
  warnings.warn(*args, **kwargs)
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/utilities/distributed.py:37: RuntimeWarning: Can save best model only with loss available, skipping.
  warnings.warn(*args, **kwargs)
```
If I print the elements of `result`, I get `{'val_loss': tensor(0.6092, device='cuda:0')}`, so `val_loss` is a tensor already.

My last attempt was to define my callback:
```
from pytorch_lightning.callbacks import ModelCheckpoint
checkpoint_callback = ModelCheckpoint(
    monitor='val_loss',
    mode='min',
)

trainer = pl.Trainer(..., checkpoint_callback=checkpoint_callback)
```
but unfortunately I still got:
```
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/utilities/distributed.py:37: RuntimeWarning: The metric you returned None must be a `torch.Tensor` instance, checkpoint not saved HINT: what is the value of val_loss in validation_epoch_end()?
  warnings.warn(*args, **kwargs)
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/utilities/distributed.py:37: RuntimeWarning: Can save best model only with val_loss available, skipping.
  warnings.warn(*args, **kwargs)
```
Would you be able to help?

### Expected behavior

I would expect Pytorch-lightning to work with minimal boilerplate (e.g. only with, `training_step`, `validation_step`). Perhaps, checkpoints should be disabled in `Trainer()` by default. Or, `validation_step()` could return automatically the aggregated metric by default. Otherwise, we may need a better explanation on how to get the aggregated metrics out the validation step, and make more clear that `validation_epoch_end()` is needed 👍 

### Environment
I am using Google Colab:
```
* CUDA:
	- GPU:
		- Tesla K80
	- available:         True
	- version:           10.1
* Packages:
	- numpy:             1.18.5
	- pyTorch_debug:     False
	- pyTorch_version:   1.6.0+cu101
	- pytorch-lightning: 0.9.1rc1
	- tensorboard:       2.2.0
	- tqdm:              4.41.1
* System:
	- OS:                Linux
	- architecture:
		- 64bit
		- 
	- processor:         x86_64
	- python:            3.6.9
	- version:           #1 SMP Thu Jul 23 08:00:38 PDT 2020
```



Discussion:
Hi! thanks for your contribution!, great first issue!
EDIT: Checkpoints work if I set them in the `EvalResult` like:
```
    def validation_step(self, batch, batch_idx):
      x, y = batch
      y_hat = self(x)
      loss = F.cross_entropy(y_hat, y)
      result = pl.EvalResult(checkpoint_on=loss)
      result.log('val_loss', loss, prog_bar=True, on_step=False, on_epoch=True)
      return result
```
but 1) I would like to keep callbacks separate from the module and only in the `Trainer()`. 2) Will setting the checkpoint in `EvalResult` interfere with other callbacks set in `Trainer()`? 3) How can I set up callbacks to monitor custom (aggregate) metrics?
@JackCaster mind sending a PR with this failing example as minimal test so we can fix it and also prevent this case in the future... 🐰 
> @JackCaster mind sending a PR with this failing example as minimal test so we can fix it and also prevent this case in the future... 🐰 

Surely I can do that. I am on vacation atm without pc. I will send it sometime next week

Resolving patch:
diff --git a/pytorch_lightning/callbacks/model_checkpoint.py b/pytorch_lightning/callbacks/model_checkpoint.py
--- a/pytorch_lightning/callbacks/model_checkpoint.py
+++ b/pytorch_lightning/callbacks/model_checkpoint.py
@@ -22,6 +22,7 @@
 
 import os
 import re
+from copy import deepcopy
 from typing import Any, Dict, Optional
 
 import numpy as np
@@ -219,7 +220,8 @@ def check_monitor_top_k(self, current) -> bool:
 
         monitor_op = {"min": torch.lt, "max": torch.gt}[self.mode]
 
-        return monitor_op(current, self.best_k_models[self.kth_best_model_path])
+        val = monitor_op(current, self.best_k_models[self.kth_best_model_path])
+        return val
 
     @classmethod
     def _format_checkpoint_name(
@@ -351,7 +353,11 @@ def on_validation_end(self, trainer, pl_module):
 
         self.epoch_last_check = epoch
 
-        ckpt_name_metrics = trainer.logger_connector.logged_metrics
+        # anything logged or in callbacks can be in the name
+        ckpt_name_metrics = deepcopy(trainer.logger_connector.logged_metrics)
+        ckpt_name_metrics.update(trainer.logger_connector.callback_metrics)
+        ckpt_name_metrics.update(trainer.logger_connector.progress_bar_metrics)
+
         filepath = self.format_checkpoint_name(epoch, ckpt_name_metrics)
         version_cnt = 0
         while self._fs.exists(filepath):
@@ -366,18 +372,19 @@ def on_validation_end(self, trainer, pl_module):
 
             if not isinstance(current, torch.Tensor):
                 rank_zero_warn(
-                    f"The metric you returned {current} must be a `torch.Tensor` instance, checkpoint not saved"
-                    f" HINT: what is the value of {self.monitor} in validation_epoch_end()?",
+                    f"The metric you returned {self.monitor}={current} must be a `torch.Tensor` "
+                    f"instance, checkpoint not saved HINT: what is the value of {self.monitor}?",
                     RuntimeWarning,
                 )
                 if current is not None:
-                    current = torch.tensor(current)
+                    current = torch.tensor(current).to(pl_module.device)
 
             if current is None:
-                rank_zero_warn(
-                    f"Can save best model only with {self.monitor} available, skipping.",
-                    RuntimeWarning,
-                )
+                m = f"Can save best model only with {self.monitor} available, skipping."
+                if self.monitor == 'checkpoint_on':
+                    m = f'No checkpoint_on found. Hint: Did you set it in EvalResult(checkpoint_on=tensor) or ' \
+                        f'TrainResult(checkpoint_on=tensor)?'
+                rank_zero_warn(m, RuntimeWarning)
             elif self.check_monitor_top_k(current):
                 self._do_check_save(filepath, current, epoch, trainer, pl_module)
             elif self.verbose:
diff --git a/pytorch_lightning/trainer/connectors/logger_connector.py b/pytorch_lightning/trainer/connectors/logger_connector.py
--- a/pytorch_lightning/trainer/connectors/logger_connector.py
+++ b/pytorch_lightning/trainer/connectors/logger_connector.py
@@ -172,8 +172,10 @@ def __log_evaluation_epoch_metrics_2(self, eval_results, test_mode):
                 # log metrics
                 self.trainer.logger_connector.log_metrics(log_metrics, {})
 
-                # track metrics for callbacks
+                # track metrics for callbacks (all prog bar, logged and callback metrics)
                 self.trainer.logger_connector.callback_metrics.update(callback_metrics)
+                self.trainer.logger_connector.callback_metrics.update(log_metrics)
+                self.trainer.logger_connector.callback_metrics.update(prog_bar_metrics)
 
                 if len(dataloader_result_metrics) > 0:
                     eval_loop_results.append(dataloader_result_metrics)
@@ -263,16 +265,18 @@ def log_train_epoch_end_metrics(self,
         # --------------------------
         # track results
         # --------------------------
-        # add the metrics to the loggers
+        # add the metrics to the loggers and callbacks
         if epoch_log_metrics and len(epoch_log_metrics) > 0:
             self.log_metrics(epoch_log_metrics, {})
+            self.callback_metrics.update(epoch_log_metrics)
 
         # add metrics to callbacks
         self.callback_metrics.update(epoch_callback_metrics)
 
-        # add metrics to progress_bar
+        # add metrics to progress_bar and callbacks
         if len(epoch_progress_bar_metrics) > 0:
             self.add_progress_bar_metrics(epoch_progress_bar_metrics)
+            self.callback_metrics.update(epoch_progress_bar_metrics)
 
     def __auto_reduce_results_on_epoch_end(self, epoch_output):
         epoch_log_metrics = {}
@@ -326,3 +330,4 @@ def save_train_loop_metrics_to_loggers(self, batch_idx, batch_output):
             grad_norm_dic = batch_output.grad_norm_dic
             if len(metrics) > 0 or len(grad_norm_dic) > 0:
                 self.log_metrics(metrics, grad_norm_dic)
+                self.callback_metrics.update(metrics)
diff --git a/pytorch_lightning/trainer/logging.py b/pytorch_lightning/trainer/logging.py
--- a/pytorch_lightning/trainer/logging.py
+++ b/pytorch_lightning/trainer/logging.py
@@ -146,6 +146,8 @@ def process_dict_result(self, output, train=False):
         # detach all metrics for callbacks to prevent memory leaks
         # no .item() because it will slow things down
         callback_metrics = recursive_detach(callback_metrics)
+        progress_bar_metrics = recursive_detach(progress_bar_metrics)
+        log_metrics = recursive_detach(log_metrics)
 
         # replace loss with checkpoint_on
         if 'loss' in callback_metrics:


Test patch:
diff --git a/tests/trainer/test_eval_loop_dict_return.py b/tests/trainer/test_eval_loop_dict_return.py
--- a/tests/trainer/test_eval_loop_dict_return.py
+++ b/tests/trainer/test_eval_loop_dict_return.py
@@ -136,7 +136,7 @@ def test_validation_step_dict_return(tmpdir):
         assert k in eval_results[1]
 
     # ensure all the keys ended up as candidates for callbacks
-    assert len(trainer.logger_connector.callback_metrics) == 8
+    assert len(trainer.logger_connector.callback_metrics) in [8, 9]
 
     # make sure correct steps were called
     assert model.validation_step_called
@@ -211,7 +211,7 @@ def test_val_step_step_end(tmpdir):
         assert k in eval_results[1]
 
     # ensure all the keys ended up as candidates for callbacks
-    assert len(trainer.logger_connector.callback_metrics) == 9
+    assert len(trainer.logger_connector.callback_metrics) in [9, 10]
 
     # make sure correct steps were called
     assert model.validation_step_called
@@ -254,7 +254,7 @@ def test_no_val_step_end(tmpdir):
         assert k in eval_results
 
     # ensure all the keys ended up as candidates for callbacks
-    assert len(trainer.logger_connector.callback_metrics) == 9
+    assert len(trainer.logger_connector.callback_metrics) in [9, 10]
 
     # make sure correct steps were called
     assert model.validation_step_called
@@ -297,7 +297,7 @@ def test_full_val_loop(tmpdir):
         assert k in eval_results
 
     # ensure all the keys ended up as candidates for callbacks
-    assert len(trainer.logger_connector.callback_metrics) == 10
+    assert len(trainer.logger_connector.callback_metrics) in [10, 11]
 
     # make sure correct steps were called
     assert model.validation_step_called
diff --git a/tests/trainer/test_trainer_steps_result_return.py b/tests/trainer/test_trainer_steps_result_return.py
--- a/tests/trainer/test_trainer_steps_result_return.py
+++ b/tests/trainer/test_trainer_steps_result_return.py
@@ -58,6 +58,7 @@ def test_training_step_result_log_step_only(tmpdir):
         assert len(logged_metrics) == 4
 
     # make sure we are using the correct metrics for callbacks
+    assert len(trainer.logger_connector.callback_metrics) == 8
     assert trainer.logger_connector.callback_metrics['checkpoint_on'] == 171
 
     # make sure pbar metrics are correct ang log metrics did not leak
@@ -123,6 +124,8 @@ def test_training_step_result_log_epoch_only(tmpdir):
     assert not model.training_step_end_called
     assert not model.training_epoch_end_called
 
+    assert len(trainer.logger_connector.callback_metrics) == 12
+
     # make sure correct metrics are logged (one per batch step as requested)
     assert len(trainer.dev_debugger.logged_metrics) == epochs
     epoch_metrics = trainer.dev_debugger.logged_metrics
@@ -198,6 +201,8 @@ def test_training_step_result_log_step_and_epoch(tmpdir):
     assert not model.training_step_end_called
     assert not model.training_epoch_end_called
 
+    assert len(trainer.logger_connector.callback_metrics) == 8
+
     # make sure correct metrics are logged (one per batch step as requested)
     assert len(trainer.dev_debugger.logged_metrics) == (epochs * batches) + epochs
     epoch_metrics = trainer.dev_debugger.logged_metrics
@@ -323,6 +328,8 @@ def test_training_step_epoch_end_result(tmpdir):
     )
     trainer.fit(model)
 
+    assert len(trainer.logger_connector.callback_metrics) == 11
+
     # make sure correct steps were called
     assert model.training_step_called
     assert not model.training_step_end_called
@@ -403,6 +410,8 @@ def test_no_auto_callbacks_with_train_loop_only(tmpdir):
     )
     trainer.fit(model)
 
+    assert len(trainer.logger_connector.callback_metrics) == 2
+
     all_losses = trainer.dev_debugger.saved_train_losses
     assert len(all_losses) == batches * epochs
 
diff --git a/tests/trainer/test_validation_steps_result_return.py b/tests/trainer/test_validation_steps_result_return.py
--- a/tests/trainer/test_validation_steps_result_return.py
+++ b/tests/trainer/test_validation_steps_result_return.py
@@ -278,6 +278,8 @@ def test_val_step_epoch_step_metrics(tmpdir):
     )
     trainer.fit(model)
 
+    assert len(trainer.logger_connector.callback_metrics) == 7
+
     # make sure correct steps were called
     assert model.validation_step_called
     assert not model.validation_step_end_called
@@ -352,6 +354,8 @@ def test_val_step_epoch_end_result(tmpdir):
     )
     trainer.fit(model)
 
+    assert len(trainer.logger_connector.callback_metrics) == 6
+
     # make sure correct steps were called
     assert model.validation_step_called
     assert not model.validation_step_end_called