checkpoints

Read and write extra metadata in PyTorch checkpoints files. These files are zip archives containing the model weights.

anemoi.utils.checkpoints.has_metadata(path: str, *, name: str = 'anemoi.json') → bool

Check if a checkpoint file has a metadata file.

Parameters:

• path (str) – The path to the checkpoint file

• name (str, optional) – The name of the metadata file in the zip archive

Returns:

True if the metadata file is found

Return type:

bool

anemoi.utils.checkpoints.get_metadata_path(path: str, *, name: str = 'anemoi.json') → str

Get the full path of the metadata file in the checkpoint.

Parameters:

• path (str) – The path to the checkpoint file

• name (str, optional) – The name of the metadata file in the zip archive

Returns:

The full path of the metadata file in the zip archive

Return type:

str

Raises:

• FileNotFoundError – If the metadata file is not found

• ValueError – If multiple metadata files are found

anemoi.utils.checkpoints.load_metadata(path: str, *, supporting_arrays: Literal[False] = False, name: str = DEFAULT_NAME) → dict

anemoi.utils.checkpoints.load_metadata(path: str, *, supporting_arrays: Literal[True] = True, name: str = DEFAULT_NAME) → tuple[dict, dict]

Load metadata from a checkpoint file.

Parameters:

• path (str) – The path to the checkpoint file

• supporting_arrays (bool, optional) – If True, the function will return a dictionary with the supporting arrays

• name (str, optional) – The name of the metadata file in the zip archive

Returns:

The content of the metadata file from JSON

Return type:

dict

Raises:

ValueError – If the metadata file is not found

anemoi.utils.checkpoints.load_supporting_arrays(zipf: ZipFile, entries: dict) → dict

Load supporting arrays from a zip file.

Parameters:

• zipf (zipfile.ZipFile) – The zip file

• entries (dict) – A dictionary of entries with paths, shapes, and dtypes

Returns:

A dictionary of supporting arrays

Return type:

dict

anemoi.utils.checkpoints.save_metadata(path: str, metadata: dict, *, supporting_arrays: dict | None = None, name: str = 'anemoi.json', folder: str = 'anemoi-metadata') → None

Save metadata to a checkpoint file.

Parameters:

• path (str) – The path to the checkpoint file

• metadata (dict) – A JSON serializable object

• supporting_arrays (dict | None, optional) – A dictionary of supporting NumPy arrays

• name (str, optional) – The name of the metadata file in the zip archive

• folder (str, optional) – The folder where the metadata file will be saved

anemoi.utils.checkpoints.replace_metadata(path: str, metadata: dict, supporting_arrays: dict | None = None, *, name: str = 'anemoi.json') → None

Replace metadata in a checkpoint file.

Parameters:

• path (str) – The path to the checkpoint file

• metadata (dict) – A JSON serializable object

• supporting_arrays (dict, optional) – A dictionary of supporting NumPy arrays

• name (str, optional) – The name of the metadata file in the zip archive

anemoi.utils.checkpoints.remove_metadata(path: str, *, name: str = 'anemoi.json') → None

Remove metadata from a checkpoint file.

Parameters:

• path (str) – The path to the checkpoint file

• name (str, optional) – The name of the metadata file in the zip archive