How do you save your model's code?

Hi all,
Coming from Keras (TensorFlow), I used to have a somewhat robust way to save/load models trained in previous development iterations. Now using PyTorch, I learned that it’s preferable to save only the model parameters contained in the model.state_dict(). I would like to be able to edit my model code, retrain the model and then perform various experiments on this last model and also its predecessors. I can only think of two naive approaches:

  1. Have a model definition function that supports a version argument. Save the model version with the model state dict and, during loading, use it to build the good model.

pros: explicit
cons: would get messy pretty quickly, would require lots of discipline to keep the old code unchanged

  1. Save the code of the model with the model state dict. During loading, dynamically use this code to build the model.

pros: less messy
cons: would probably get complex if the model definition function and its dependencies don’t fit in a single file (which, excluding libraries, is not my case for now)

Cleary both approaches could still break if I change my python environment (update the project dependencies) but I guess I can live with that. Do you think these approaches are viable? Do you know better?

I don’t think it is possible to save raw written code. Torchscript can be used, where the model is converted to torchscript and that can be saved.

Thank you. I will definitely look into torchscript eventually. For now, I just save the code of my model building function as a string (obtains with the inspect module). When loading, I use exec. This is probably not a really clean way to serialize a function but it works for my needs. The dill package might be a better way. I also realized my question is really specific to the way I work, where each revision only defines one version of the model.

1 Like