Mentions légales du service

Skip to content

Improve portability of .pth files

E. Madison Bray requested to merge embray/issue-53 into master

This is a workaround for #53 (closed). I'm not sure if it's the best solution long-term but it should fix some of the major pain points:

  • It does not try to load configuration information that is only relevant to the machine the model was trained on.

  • It does not use configuration information that is only relevant to the machine the model was trained on (e.g. options related to CUDA devices, etc.)

For now this works around the issue of information about paths in the config file by simply not including the entire training config in the .pth file anymore. I have mixed feelings about that, and it's something we could revisit in the future. But for now the .pth file and the actual training config can live side-by-side.

Instead, the .pth file contains only portions of the configuration required to load the network and run predictions on it:

  • network: The name and init params for the network.
  • preprocessing: Not strictly required, but the dnadna predict --preprocess allows pre-processing the input data with the same options as the training data while running predictions.
  • learned_params: For prediction we need to know what parameters the model can predict about the dataset.
  • train_mean / train_std: Needed to de-standardize predictions (!62 (merged))
  • state_dict: The actual parameters of the network dumped from pytorch.
  • dnadna_version / model_datetime: Software version and timestamp. We will probably want to include additional metadata in the future, as well as information about plugins used (#71)

I think this is "good enough" for version 1, but for the future we will want to have a meeting to really pin down the format of our files and standardize it better. This will help, also, to implement backwards-compatibility so that as much as possible newer versions of DNADNA will be able to load models trained with older versions.

Merge request reports