[Feature] Transfer Learning

[icedoom888]

Is your feature request related to a problem? Please describe.

Implementation of the a Transfer Learning feature: initialise the model state from a previous experiment.

Necessary to reproduce: Regional data-driven weather modeling with a global stretched-grid

Describe the solution you'd like

• Sanify checkpoint when size missmatch: remove all tensors that do not fit.
• Save as new checkpoint suited for Loading.
• Load the new checkpoint with strict=False,
• Start a new training run.

Describe alternatives you've considered

No response

Additional context

No response

Organisation

MeteoSwiss

[gabrieloks]

Hello @icedoom888. Thanks a lot for opening this issue. We have already started working on the first phase of transfer learning (going from stage A to stage B in your figure) in this PR #140 . I have some questions about the solution:

-When there are size mismatches, do we really need to remove these tensors? We could just skip loading them and have the model initialise them at their default value.

-Do we need to save the sanified checkpoint and load it again? Could we just load the checkpoint from the previous stage and avoid loading certain layers with a filter?

Happy to have a chat with you if needed.

Gabriel

[icedoom888]

Hey @gabrieloks!
I was able to run all stages with my implementation, will just need to polish it up a bit ;)

Regarding your questions: I am addressing this issue right now, will try to load the model directly without storing the transfer.ckpt file.
I will update you as soon as i test the new feature :)