Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Updating calculator weights from checkpoints #845

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

LarsSchaaf
Copy link
Collaborator

If you want to see performance changes throughout epochs, you can load a single model and then update the weights from former checkpoints.

To do:

  • create test: where you compare the output from loading two separate models or updating based on their corresponding checkpoints.

@RokasEl
Copy link
Collaborator

RokasEl commented Mar 6, 2025

What's the use case of doing:

calc.update_weights_from_checkpoint(LOAD_PATH)

compared to:

calc = MACECalculator(model_paths=LOAD_PATH)

?

If there is a genuine case, I think it would be better to modularize the __init__ of the calculator to use a load_weights function and use that for both initializing the calculator, and updating the weights inplace.

@LarsSchaaf
Copy link
Collaborator Author

LarsSchaaf commented Mar 6, 2025

The difference is that one loads .model s and this one can load checkpoints .pt files, ie just state_dicts

        Example:
            >>> atoms = read('path/to/atoms_file.xyz')
            >>> atoms.calc = MACECalculator('path/to/mace_model.model')
            >>> print("Energy before:", atoms.get_potential_energy())
            >>> atoms.calc.update_weights_from_checkpoint('path/to/checkpoint.pt')
            >>> print("Energy after:", atoms.get_potential_energy())

@LarsSchaaf
Copy link
Collaborator Author

Note:

  • state_dict changes when using cuEquivariance, so can only update models with checkpoints with same cuequivariance setting
  • compiled models also don't work atm

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants