Repository: Lightning-AI/pytorch-lightning

Issue:
training_epoch_end's outputs doesn't have 'loss' key
pytorch-lightning: build from master
```
Traceback (most recent call last):
  File "main.py", line 140, in <module>
    main(hparams)
  File "main.py", line 72, in main
    trainer.fit(model)
  File "/mnt/lustre/maxiao1/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 881, in fit
    self.ddp_train(task, model)
  File "/mnt/lustre/maxiao1/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/distrib_data_parallel.py", line 539, in ddp_train
    self.run_pretrain_routine(model)
  File "/mnt/lustre/maxiao1/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1091, in run_pretrain_routine
    self.train()
  File "/mnt/lustre/maxiao1/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 376, in train
    self.run_training_epoch()
  File "/mnt/lustre/maxiao1/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 510, in run_training_epoch
    self.run_training_epoch_end(epoch_output)
  File "/mnt/lustre/maxiao1/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 535, in run_training_epoch_end
    epoch_output = model.training_epoch_end(epoch_output)
  File "/mnt/lustre/maxiao1/PVM/models/baseline.py", line 335, in training_epoch_end
    avg_loss = torch.stack([x['loss'] for x in outputs]).mean()
  File "/mnt/lustre/maxiao1/PVM/models/baseline.py", line 335, in <listcomp>
    avg_loss = torch.stack([x['loss'] for x in outputs]).mean()
KeyError: 'loss'
```

This is my code:
```
    def training_step(self, batch, batch_idx):
        ...
        return {'loss': loss, "train_acc": acc}

    def training_epoch_end(self, outputs):
        avg_loss = torch.stack([x['loss'] for x in outputs]).mean()
        avg_acc = torch.stack([x['train_acc'] for x in outputs]).mean()
        logs = {'loss': avg_loss, 'train_acc': avg_acc}
        progress_bar = {'train_loss': avg_loss, 'train_acc': avg_acc}
        results = {
            'log': logs,
            'progress_bar': progress_bar
        }
        return results
```


Discussion:
Try: `avg_loss = torch.stack([x['batch_loss'] for x in outputs]).mean()`
Thanks， it works
but 'train_acc' key doesn't exist, neither do `batch_train_acc`. How to access other keys returned in training_step?
As of now in lightning you can access them using `x['callback_metrics']['loss']` and `x['callback_metrics']['train_acc']`, but I think it should be handled in a similar way we do this with `validation_epoch_end` and `test_epoch_end`.
Hi! One hint: for me it works with "loss" under windows but not under ubuntu.
Weird!! Why is this think platform dependent?? :thinking: 
@Pet222 , are u sure that versions on ubuntu and windows are same?
Hey @williamFalcon is this intended behaviour? I was surprised to see this breaking change being introduced with no warning. 
If it is intended, why not have consistent behaviour over `validation_epoch_end` and `test_epoch_end`.

If it is not intended, as it seems due to the "bug fix" tag, are you working on it or should I make a PR for this?
what is the behavior? that the "loss" key is not in training_epoch_end? If so, that's a bug because it should be there
@williamFalcon , on the latest version, the `loss` key was changed to the `batch_loss`. I think it was changed [here](https://github.com/PyTorchLightning/pytorch-lightning/commit/0f073819d3e0df8db7602eab489b1bad0fc0949c#diff-c45bd21c331565cbe62aaa12fa43aa0aR717)
Yes, the fact that you need to access it through 'callback metrics'.
Got it!

On Tue, 30 Jun 2020 at 12:44, William Falcon <notifications@github.com>
wrote:

> what is the behavior? that the "loss" key is not in training_epoch_end? If
> so, that's a bug because it should be there
>
> —
> You are receiving this because you commented.
> Reply to this email directly, view it on GitHub
> <https://github.com/PyTorchLightning/pytorch-lightning/issues/2372#issuecomment-651740702>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/ABKWP6XTUJDTEDJ2NZQ3RKTRZHFY5ANCNFSM4OJKX4KQ>
> .
>
-- 
Best Regards,
Miguel Vera

+351 915 198 452
miguel.coimbra.vera@protonmail.com
Github/Captainvera <http://www.github.com/captainvera>

@captainvera would love a PR :) 

Resolving patch:
diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
--- a/pytorch_lightning/trainer/training_loop.py
+++ b/pytorch_lightning/trainer/training_loop.py
@@ -776,6 +776,7 @@ def optimizer_closure(self, split_batch, batch_idx, opt_idx, optimizer, hiddens)
             # PROCESS THE RESULT
             # ----------------------------
             # format and reduce outputs accordingly
+            training_step_output_for_epoch_end = training_step_output
             training_step_output = self.process_output(training_step_output, train=True)
 
             # TODO: temporary part of structured results PR
@@ -788,7 +789,7 @@ def optimizer_closure(self, split_batch, batch_idx, opt_idx, optimizer, hiddens)
             )
 
             # if the user decides to finally reduce things in epoch_end, save raw output without graphs
-            training_step_output_for_epoch_end = recursive_detach(training_step_output)
+            training_step_output_for_epoch_end = recursive_detach(training_step_output_for_epoch_end)
 
         # accumulate loss
         # (if accumulate_grad_batches = 1 no effect)


Test patch:
diff --git a/tests/base/deterministic_model.py b/tests/base/deterministic_model.py
--- a/tests/base/deterministic_model.py
+++ b/tests/base/deterministic_model.py
@@ -58,7 +58,7 @@ def training_step_dict_return(self, batch, batch_idx):
         pbar = {'pbar_acc1': torch.tensor(17).type_as(acc), 'pbar_acc2': torch.tensor(19).type_as(acc)}
 
         self.training_step_called = True
-        return {'loss': acc, 'log': logs, 'progress_bar': pbar}
+        return {'loss': acc, 'log': logs, 'progress_bar': pbar, 'train_step_test': torch.tensor(549).type_as(acc)}
 
     def training_step_for_step_end_dict(self, batch, batch_idx):
         """sends outputs to training_batch_end"""
@@ -89,11 +89,11 @@ def training_step_end_dict(self, output):
         assert 'pbar_acc1' in output
         assert 'pbar_acc2' in output
 
-        logs = {'log_acc1': output['log_acc1'], 'log_acc2': output['log_acc2']}
-        pbar = {'pbar_acc1': output['pbar_acc1'], 'pbar_acc2': output['pbar_acc2']}
+        logs = {'log_acc1': output['log_acc1'] + 2, 'log_acc2': output['log_acc2'] + 2}
+        pbar = {'pbar_acc1': output['pbar_acc1'] + 2, 'pbar_acc2': output['pbar_acc2'] + 2}
 
         acc = output['loss']
-        return {'loss': acc, 'log': logs, 'progress_bar': pbar}
+        return {'loss': acc, 'log': logs, 'progress_bar': pbar, 'train_step_end': acc}
 
     def training_epoch_end_dict(self, outputs):
         self.training_epoch_end_called = True
@@ -104,12 +104,14 @@ def training_epoch_end_dict(self, outputs):
             # only saw 4 batches
             assert len(outputs) == 4
             for batch_out in outputs:
-                assert len(batch_out.keys()) == 5
-                keys = ['batch_loss', 'pbar_on_batch_end', 'log_metrics', 'callback_metrics']
+                assert len(batch_out.keys()) == 4
+                assert self.count_num_graphs(batch_out) == 0
+                last_key = 'train_step_end' if self.training_step_end_called else 'train_step_test'
+                keys = ['loss', 'log', 'progress_bar', last_key]
                 for key in keys:
                     assert key in batch_out
 
-        prototype_loss = outputs[0]['batch_loss']
+        prototype_loss = outputs[0]['loss']
         logs = {'epoch_end_log_1': torch.tensor(178).type_as(prototype_loss)}
         pbar = {'epoch_end_pbar_1': torch.tensor(234).type_as(prototype_loss)}
 
diff --git a/tests/trainer/test_trainer_steps.py b/tests/trainer/test_trainer_steps.py
--- a/tests/trainer/test_trainer_steps.py
+++ b/tests/trainer/test_trainer_steps.py
@@ -31,7 +31,12 @@ def test_training_step_dict(tmpdir):
     assert out.batch_log_metrics['log_acc1'] == 12.0
     assert out.batch_log_metrics['log_acc2'] == 7.0
 
-    pbar_metrics = out.training_step_output_for_epoch_end['pbar_on_batch_end']
+    train_step_out = out.training_step_output_for_epoch_end
+    pbar_metrics = train_step_out['progress_bar']
+    assert 'loss' in train_step_out
+    assert 'log' in train_step_out
+    assert 'progress_bar' in train_step_out
+    assert train_step_out['train_step_test'] == 549
     assert pbar_metrics['pbar_acc1'] == 17.0
     assert pbar_metrics['pbar_acc2'] == 19.0
 
@@ -59,12 +64,14 @@ def training_step_with_step_end(tmpdir):
 
     out = trainer.run_training_batch(batch, batch_idx)
     assert out.signal == 0
-    assert out.batch_log_metrics['log_acc1'] == 12.0
-    assert out.batch_log_metrics['log_acc2'] == 7.0
+    assert out.batch_log_metrics['log_acc1'] == 14.0
+    assert out.batch_log_metrics['log_acc2'] == 9.0
 
-    pbar_metrics = out.training_step_output_for_epoch_end['pbar_on_batch_end']
-    assert pbar_metrics['pbar_acc1'] == 17.0
-    assert pbar_metrics['pbar_acc2'] == 19.0
+    train_step_end_out = out.training_step_output_for_epoch_end
+    pbar_metrics = train_step_end_out['progress_bar']
+    assert 'train_step_end' in train_step_end_out
+    assert pbar_metrics['pbar_acc1'] == 19.0
+    assert pbar_metrics['pbar_acc2'] == 21.0
 
 
 def test_full_training_loop_dict(tmpdir):
@@ -99,12 +106,13 @@ def test_full_training_loop_dict(tmpdir):
 
     out = trainer.run_training_batch(batch, batch_idx)
     assert out.signal == 0
-    assert out.batch_log_metrics['log_acc1'] == 12.0
-    assert out.batch_log_metrics['log_acc2'] == 7.0
+    assert out.batch_log_metrics['log_acc1'] == 14.0
+    assert out.batch_log_metrics['log_acc2'] == 9.0
 
-    pbar_metrics = out.training_step_output_for_epoch_end['pbar_on_batch_end']
-    assert pbar_metrics['pbar_acc1'] == 17.0
-    assert pbar_metrics['pbar_acc2'] == 19.0
+    train_step_end_out = out.training_step_output_for_epoch_end
+    pbar_metrics = train_step_end_out['progress_bar']
+    assert pbar_metrics['pbar_acc1'] == 19.0
+    assert pbar_metrics['pbar_acc2'] == 21.0
 
 
 def test_train_step_epoch_end(tmpdir):
@@ -138,6 +146,7 @@ def test_train_step_epoch_end(tmpdir):
     assert out.batch_log_metrics['log_acc1'] == 12.0
     assert out.batch_log_metrics['log_acc2'] == 7.0
 
-    pbar_metrics = out.training_step_output_for_epoch_end['pbar_on_batch_end']
+    train_step_end_out = out.training_step_output_for_epoch_end
+    pbar_metrics = train_step_end_out['progress_bar']
     assert pbar_metrics['pbar_acc1'] == 17.0
     assert pbar_metrics['pbar_acc2'] == 19.0