MACE Training Output Models

[jasonlu38]

Hello! I'm wondering about the results for training a MACE model, usually I see "_stage_one.model" and "_stage_two.model" along with when I specify a seed for the training I see (for example 123) "_run-123_stage_one.model" and "_run-123_stage_two.model". I also see just "*.model" of whatever name I specified for the model in the MACE training.

I'm wondering which model would be the "final" output model since I do see the Colab notebooks use the "*_stage_two.model"? Let me know as I'd love to know the specifics and am having a hard time understanding the output of the logs (which models are which)!
For fine-tuning, it wasn't obvious to me what the difference between the default and pt_heads were. Is the pt_head the finetuned model, while the default model would be the same model as if I didn't fine-tune? Thank you very much!

[ilyes319]

Hey @jasonlu38,

The stage stage_one.model and stage_two.model corresponds to the models saved at the two stages of training of MACE, where the first one trains mostly on forces and second one more on the energies, as we found that procedure to help converge the energies errors substantially. At the end of the log file, you can see the statistics of the two stages of the model, and pick the one that has the best performance for your needs.

I recommend using the model that is saved in checkpoint folder. Does not use the model files with "compiled" in them as these are torchscript models that are meant to be use for (non-MLIAP) lammps.

For the fine-tuning, the pt_head is the pre-trained head and default the fine-tuned head if you do not change the name of it. You can have more info here: https://mace-docs.readthedocs.io/en/latest/guide/multihead_finetuning.html