core.scripts.migrate_checkpoint

core.scripts.migrate_checkpoint#

Copyright (c) Meta Platforms, Inc. and affiliates.

This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Attributes#

parser

Classes#

RenameUnpickler

This takes a binary file for reading a pickle data stream.

Functions#

find_new_module_name(module)

update_config(config_or_data)

generate_stress_task_config(dataset_name, task_name, rmsd)

migrate_checkpoint(→ dict)

Migrates a checkpoint by updating module imports and configurations.

Module Contents#

core.scripts.migrate_checkpoint.find_new_module_name(module)#
core.scripts.migrate_checkpoint.update_config(config_or_data)#
class core.scripts.migrate_checkpoint.RenameUnpickler(file, *args, **kwargs)#

Bases: pickle.Unpickler

This takes a binary file for reading a pickle data stream.

The protocol version of the pickle is detected automatically, so no protocol argument is needed. Bytes past the pickled object’s representation are ignored.

The argument file must have two methods, a read() method that takes an integer argument, and a readline() method that requires no arguments. Both methods should return bytes. Thus file can be a binary file object opened for reading, an io.BytesIO object, or any other custom object that meets this interface.

Optional keyword arguments are fix_imports, encoding and errors, which are used to control compatibility support for pickle stream generated by Python 2. If fix_imports is True, pickle will try to map the old Python 2 names to the new names used in Python 3. The encoding and errors tell pickle how to decode 8-bit string instances pickled by Python 2; these default to ‘ASCII’ and ‘strict’, respectively. The encoding can be ‘bytes’ to read these 8-bit string instances as bytes objects.

find_class(module, name)#

Return an object from a specified module.

If necessary, the module will be imported. Subclasses may override this method (e.g. to restrict unpickling of arbitrary classes and functions).

This method is called whenever a class or a function object is needed. Both arguments passed are str objects.

core.scripts.migrate_checkpoint.generate_stress_task_config(dataset_name, task_name, rmsd)#
core.scripts.migrate_checkpoint.migrate_checkpoint(checkpoint_path: torch.nn.Module, rm_static_keys: bool = True, map_undefined_stress_to: str | None = None, add_stress: bool = False, task_add_stress: str | None = None, model_version: float = 1.0) → dict#

Migrates a checkpoint by updating module imports and configurations.

This function loads a checkpoint, updates its configuration using the mapping defined in fairchem.core.scripts.migrate_imports,

optionally adds stress tasks for datasets that don’t have them, and optionally removes static keys that are no longer needed.

Parameters:

  • checkpoint_path – Path to the input checkpoint file

  • rm_static_keys – Whether to remove static keys from the state dictionaries

  • task_add_stress – If provided, adds stress tasks for datasets based on this task

  • model_version – Inject this model version into model

Returns:

Migrated checkpoint dict

core.scripts.migrate_checkpoint.parser#
Contents