Skip to content

Create a robust method for saving/loading models

Created by: RasmusOrsoe

Is your feature request related to a problem? Please describe. Loading models saved in a different version of the repo, or under a different version of torch/torch geometric can be a real pain in the butt.

Describe the solution you'd like

  • Make the Model contain information about the environment in which it was created. E.g. Model.environment. The field should track version control of the sensitive dependencies; graphnet (hash commit), torch, torch-geometric etc. Not pandas, numpy..
  • In line with #262 , create a saving method in Model that saves the state_dict and a list of components (gnn, detector, task etc) when training is complete.
  • Create a utility function build_model(state_dict, components) that accepts the state_dict and components and returns a full Model. This method should be able to handle name changes, moving of code etc.

Additional context I think this is necessary for backwards-compatibility. We need to be able to run "old" models from "old" studies. Currently this is very difficult.