feat(wandb): log models as artifacts

[borisdayma]

What does this PR do?

Fixes #4903

should be merged after #5931

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 🙃

[pep8speaks]

Hello @borisdayma! Thanks for updating this PR.

There are currently no PEP 8 issues detected in this Pull Request. Cheers! 🍻

Comment last updated at 2021-05-27 09:40:00 UTC

[codecov]

Codecov Report

Merging #6231 (e0f302f) into master (04dcb17) will decrease coverage by 5%.
The diff coverage is 100%.

@@           Coverage Diff           @@
##           master   #6231    +/-   ##
=======================================
- Coverage      92%     88%    -5%     
=======================================
  Files         199     199            
  Lines       12962   12995    +33     
=======================================
- Hits        11967   11409   -558     
- Misses        995    1586   +591     

[borisdayma]

I started from my other ongoing PR so you can see the changes here

[borisdayma]

The history looks a bit messed up so I'll try to fix it or force push from a clean master state (since there's only one real commit so far).

[borisdayma]

Here are examples:

• with default checkpointing: artifact metadata
• with custom ModelCheckpoint: artifact metadata

The idea is to avoid putting too much metadata that would hide more meaningful data.

[borisdayma]

Here is current behavior:

• default: log only last model as artifact → example artifact
• custom ModelCheckpoint: according to options (looks through last_model_path, best_model_path and best_k_models) → example artifact

We can see in the 2nd example that 3 model were logged:

• versions are ordered by timestamp
• latest and best aliases are automatically set
• file within artifact is called model.ckpt

This would let users access the best model from a specific run with a command such as:
artifact = run.use_artifact('USER/PROJECT/RUN_NAME:best', type='model')

Ideally in a future PR, it would be nice to have an option log_model='all' (in particular for long training that may crash) which probably means allowing loggers to have functionalities similar to other callbacks.

[justusschock]

I like it. Just some more comments :)

Will the artifacts be uploaded and updated continuously? E.g. will I at every point of the training have a best model artifact at wandb and if I get a new one it will replace the old one?

[tchaton]

LGTM !

[tchaton]

self.log_images could be the API used to log every image related artefacts.

[borisdayma]

Hi guys, just checking if anything else is needed from me.
I know quite a few people already who are interested in accessing this feature.

[edenlightning]

Apologize for the late reply Boris!
We want to spend some time thinking about how to make this a robust API that would work well for all supported loggers. We have 1.3 release coming soon which is too late to add this in, so we will prioritize this right after the release.

Thanks again for your contribution, and sorry we haven't been able to address quicker.

[borisdayma]

Hi, I'm just checking if there's any update or if I can do anything from my side.
It would be great to implement this feature as it will close a lot of current open issues (and make the integration so much better!)

[borisdayma]

Hey guys, just wondering if you would be ok using my current version?

Right now I just need a callback for after_save_checkpoint (which is those 2 lines + associated tests) so whenever you have a more robust one, it will be easy to switch it out.

[Borda]

@awaelchli how is it going here, I saw you had most comments so mind giving your decision? 🐰
@borisdayma mind resolve conflicts?

[borisdayma]

Ok I resolved conflicts

[borisdayma]

Hi, just checking if i need to do anything else on my side?

[borisdayma]

Hey guys, just checking on the progress here.
There's a few issues linked to this PR so I'd love to know what is still left to do.

[borisdayma]

Thanks guys!!!