ACEsuit / mace

MACE - Fast and accurate machine learning interatomic potentials with higher order equivariant message passing.
Other
415 stars 157 forks source link

Possibility to not compute energies on `MACE.forward` #338

Closed pfebrer closed 3 months ago

pfebrer commented 4 months ago

Is your feature request related to a problem? Please describe. I'm trying to think about how to use the main MACE model to train on things other than energies (i.e. having a custom readout) without having to rewrite the whole model.

Describe the solution you'd like I think having a compute_energies argument on MACE.forward that defaults to True could be great. Then one could set it to False and use the main mace model to just get node features and then do whatever you want with them. It would also be straightforward to take a model that is pretrained for energies and then start training it for something else. This argument would be minimally invasive and would not introduce any problems in terms of compatibility (I think).

Describe alternatives you've considered

ilyes319 commented 3 months ago

Hey,

MACE already returns the node features, so you can easily use it this way already. The computation of the energy adds very little cost, so that will not be a problem.