Keep only the best and the latest artifacts in wandb logger

[ohayonguy]

Hey,
I am not sure if I can currently keep only the best and latest models as a wandb artifact using the WandbLogger? That is, I am looking for a behavior similar to log_model='all', but which keeps only the latest and best models and deletes previous checkpoints from the wandb artifacts of the experiment.
My checkpoints weigh about 1GB and I don't want to keep the entire history of checkpoints with the log_model='all' flag, but rather only the best and the latest models. I thought about inheriting from the WandbLogger and following the guideline here:
https://gitbook-docs.wandb.ai/guides/artifacts/api#cleaning-up-unused-versions

Any thoughts?
Maybe this should be a feature request?

[awaelchli]

Hey @ohayonguy

Given the docs of WandbLogger I think you can just set log_model=True. This will then upload the checkpoints at the end of training. And if you have set save_top_k=n it will only upload the best n ones.

[ohayonguy]

Hey @awaelchli, thanks for the reply.

What if I don't want to wait for the end of training, and instead stop manually in the middle?
I still want to save checkpoints as artifacts during training, but without blowing up disk space (I am using wandb local).

I managed to solve this by adding some lines after https://github.com/PyTorchLightning/pytorch-lightning/blob/392c577825af1b9e2ba920f63de5569ef9b0d0cb/pytorch_lightning/loggers/wandb.py#L303

I used the suggestion in here: https://gitbook-docs.wandb.ai/guides/artifacts/api#cleaning-up-unused-versions
and it works pretty well. The artifacts directory in wandb only keeps the latest and the best checkpoints.

Do you think that this worth a PR? for instance, add a log_model=epoch option or something like that.

[tchaton]

Dear @ohayonguy,

Definitely a cool feature. Yes, worth a PR.

Best,
T.C

[ohayonguy]

@tchaton Sounds good.
#9356