[feat] Add {,load_}state_dict to ResultCollection 1/n

[tchaton]

What does this PR do?

Tracking Issue: #7898

This PR adds a mechanism to reload from state_dict and restore logged values.

Fixes #<issue_number>

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 #7948 (5c8ce17) into master (b71aa55) will decrease coverage by 0%.
The diff coverage is 97%.

@@          Coverage Diff           @@
##           master   #7948   +/-   ##
======================================
- Coverage      92%     91%   -0%     
======================================
  Files         207     207           
  Lines       13375   13464   +89     
======================================
+ Hits        12245   12295   +50     
- Misses       1130    1169   +39     

[ananthsub]

This PR adds a mechanism to reload from state_dict and restore logged values.

Could you describe more what is the problem we're trying to solve? Is this the only option for restoring logged values?

Relying on self.log to checkpoint metric states is adding even more responsibilities when I think we should be going the other direction and moving responsibilities out of self.log. #7183 (comment)

[tchaton]

This PR adds a mechanism to reload from state_dict and restore logged values.

Could you describe more what is the problem we're trying to solve? Is this the only option for restoring logged values?

Relying on self.log to checkpoint metric states is adding even more responsibilities when I think we should be going the other direction and moving responsibilities out of self.log. #7183 (comment)

Hey @ananthsub,

Currently, the ResultCollection isn't fault tolerant.

Problems:

• On failure, the collection is being lost.
• On restoration, the ResultMetric lost reference to the Metrics.

In order to resolve this, this PR adds the following:

• A state_dict and load_from_state_dict function to the ResultCollection to restore in case of failure.
• A new field in the metadata object called attribute_name. This is used to automatically re-create references to the Metric on reload by auto-inspecting the LightningModule object. If the Metric aren't part of the LightningModule (not recommended), the user should still be able to restore by providing them manually.

The restoration isn't implemented yet, it would be done in 2/n and 3/n PRs.

The state dumping / restoration will be added to CheckpointConnector. There will be a restore_result_collections function.

PR 2/n: Add ResultCollection state_dict of all loops to model checkpoint inside CheckpointConnector dump_checkpoint function

PR 3/n: Add ResultCollection restore within restoring calls.

Note: The self.log doesn't checkpoint metric states. the self.log function of the LightningModule is responsible to properly populate the Metadata by extracting the attribute_name and propagate the logged values to the current ResultCollection.
The CheckpointConnector will be responsible for this.

Best,
T.C

[tchaton]

Hey Ananth,

Actually it might be possible to get rid of the attritube name. I will try this on Monday, which will reduce self.log responsability.

Confirm I could remove it.

[awaelchli]

apart from the pickling concern LGTM
confidence low, not yet so familiar with new logger connector etc.

[SeanNaren]

Again low confidence but the plan makes sense and test for restoration looks good