Migrating saved pretrained models to new versions

PCNtoolkit saves trained models as JSON files. These files allow to reload the pretrained model parameters back to our workspace (without having to re-estimate the models - which saves a lot of time, especially for models that are trained on thousands of subjects). As PCNtoolkit develops, the internal structure of these saved files may change.

For example: a parameter may be renamed or parameters may change representation (e.g. from integer to list)

To make sure older saved models still work with newer versions of PCNtoolkit, we migrate an old saved model format into the format expected by the current version of PCNtoolkit.

Analogy: opening an old Word document

A useful analogy is opening an old .doc file in a new version of Microsoft Word. You do not manually edit the file format yourself. You simply click Open. Word detects that the file was created with an older format and automatically converts it so that the new version can read it.

PCNtoolkit does something similar. When you load an older saved model, PCNtoolkit checks which version created the model and, if needed, automatically updates the saved model information before loading it.

Every model saved by PCNtoolkit contains a field called ptk_version.

For example in both normative_model.json and regression_model.json you will find:

{
  "ptk_version": "1.2.0",
}

This tells PCNtoolkit which version of the software created the saved model.

When you load a model with

model = NormativeModel.load("path/to/saved_model")

, PCNtoolkit compares:

1. the version stored in the saved model, and
2. the version of PCNtoolkit currently installed on your computer.

What if the model was saved with an older version?

If the model was saved with an older version of PCNtoolkit and there are breaking changes, a migration function must exist in the pcntoolkit/util/migration.py that will apply the required migrations automatically.

What if the model was saved with a newer version?

If the model was saved using a newer version of PCNtoolkit than the one you currently have installed, PCNtoolkit will warn you.

For example:

This model was saved with a newer version of PCNtoolkit.
Please upgrade PCNtoolkit.

How to add a migration funciton

Step 1: Identify the component name

The components (classes) below are the ones that load data from the JSON saved model

• BLR
• HBR
• BasisFunction
• Scaler
• Likelihood
• Prior

Step 2: Write and register the migration

Open:

pcntoolkit/util/migration.py

Add the migration at the bottom of the file, under the registered migrations section.

Example: suppose you renamed the field knot_positions to knots in version 1.3.0. You need to add a migration function so that files saved before 1.3.0 can still be loaded.

To add a migration, create a new function at the bottom of pcntoolkit/util/migration.py with a @registry.register decorator that takes two arguments:

1. Which component owns the field — "BasisFunction" here, because knots lives on BasisFunction
2. When the rename happened — introduced_in="1.3.0" makes sure to only run this migration on files saved before that version.

@registry.register("BasisFunction", introduced_in="1.3.0")
def _migrate_basis_function_1_3_0(d: dict) -> dict:
    """Rename 'knot_positions' to 'knots'."""
    if "knot_positions" in d:
        d["knots"] = d.pop("knot_positions")
    return d

Step 3: Add a test

Add a test in:

test/test_util/test_migration.py

The test should check that an old dictionary is correctly converted into the new dictionary. For this test you need to add a JSON file from an old version, then load it in to the new version and check that all the contents were correctly transformed.

Rules for writing migration functions

Use the version where the new format appears

The introduced_in value must be the version where the new format first exists.

For example, if the field name changes in version 1.3.0, use:

introduced_in="1.3.0"

Do not use the old version number.

Make migrations safe to run

Always check whether the old key exists before changing it:

if "old_key" in d:
    d["new_key"] = d.pop("old_key")

This makes the migration safe even if the dictionary has already been updated.

Use a descriptive function name

Use this naming pattern:

_migrate_<component>_<version>

For example:

_migrate_basis_function_1_3_0

Version comparison

PCNtoolkit compares versions using standard Python package version rules.

For example:

1.2.0.post1 > 1.2.0
1.3.0rc1 < 1.3.0
1.3.0.dev0 < 1.3.0

Very old saved models may not contain a ptk_version field at all. For this reason PCNtoolkit treats them as if they came from version:

0.0.0

This means all relevant migrations will be applied.