Fix inconsistent outputs in on_*_end and *_end

[ethanwharris]

What does this PR do?

Fixes #6921

Several changes:

• Move training_epoch_end (which actual handles model epoch end) from logger connector to training_loop.py
• Use one static method (_prepare_outputs) in training_loop.py to unpack all Result object outputs from both epoch_end and batch_end
• Pass these unpacked outputs to both on_train_epoch_end and training_epoch_end hooks (previously, training_epoch_end behaved differently)
• Moved evaluation output handling to trainer.py so that both evaluation_epoch_end and on_evaluation_epoch_end get the same outputs argument
• Changed order of calls to training_epoch_end and on_train_epoch_end to be consistent with the docs

Previously, running the code from the issue, the user saw this output:

callback on_validation_batch_end: {'loss': tensor(1.9916), 'foo': 'from_val_step'}
validation_epoch_end: [{'loss': tensor(1.9916), 'foo': 'from_val_step'}]
callback on_validation_epoch_end: [[{'loss': tensor(1.9916), 'foo': 'from_val_step'}]]
on_validation_epoch_end: [[{'loss': tensor(1.9916), 'foo': 'from_val_step'}]]
callback on_train_batch_end: [[{'extra': {'foo': 123}, 'minimize': tensor(0.2295)}]]
callback on_train_epoch_end: [[[{'extra': {'foo': 123}, 'minimize': tensor(0.2295)}]]]
on_train_epoch_end: [[[{'extra': {'foo': 123}, 'minimize': tensor(0.2295)}]]]
training_epoch_end: [{'foo': 123, 'loss': tensor(0.2295)}]
callback on_validation_batch_end: {'loss': tensor(2.2656), 'foo': 'from_val_step'}
validation_epoch_end: [{'loss': tensor(2.2656), 'foo': 'from_val_step'}]
callback on_validation_epoch_end: [[{'loss': tensor(2.2656), 'foo': 'from_val_step'}]]
on_validation_epoch_end: [[{'loss': tensor(2.2656), 'foo': 'from_val_step'}]]

Now sees this:

callback on_validation_batch_end: {'loss': tensor(0.6422), 'foo': 'from_val_step'}
validation_epoch_end: [{'loss': tensor(0.6422), 'foo': 'from_val_step'}]
callback on_validation_epoch_end: [{'loss': tensor(0.6422), 'foo': 'from_val_step'}]
on_validation_epoch_end: [{'loss': tensor(0.6422), 'foo': 'from_val_step'}]
callback on_train_batch_end: {'foo': 123, 'loss': tensor(1.2562)}
training_epoch_end: [{'foo': 123, 'loss': tensor(1.2562)}]
callback on_train_epoch_end: [{'foo': 123, 'loss': tensor(1.2562)}]
on_train_epoch_end: [{'foo': 123, 'loss': tensor(1.2562)}]
callback on_validation_batch_end: {'loss': tensor(0.4613), 'foo': 'from_val_step'}
validation_epoch_end: [{'loss': tensor(0.4613), 'foo': 'from_val_step'}]
callback on_validation_epoch_end: [{'loss': tensor(0.4613), 'foo': 'from_val_step'}]
on_validation_epoch_end: [{'loss': tensor(0.4613), 'foo': 'from_val_step'}]

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

[codecov]

Codecov Report

Merging #6969 (d3b1735) into master (e891ceb) will decrease coverage by 0%.
The diff coverage is 98%.

@@          Coverage Diff           @@
##           master   #6969   +/-   ##
======================================
- Coverage      92%     92%   -0%     
======================================
  Files         194     194           
  Lines       12328   12329    +1     
======================================
- Hits        11324   11322    -2     
- Misses       1004    1007    +3     

[ananthsub]

@ethanwharris n00b question: should the lightning module hook come before the callback hook runs?

[ethanwharris]

@ananthsub The docs certainly think it should, so it will following this PR 😃

[ethanwharris]

@carmocca Have added some tests for training_loop 😃 - one for call order, and one checking that they are called with the correct outputs objects.

[tchaton]

LGTM !