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
Rewriting the model's code on my library using a different readout. This is not good since we need to make sure that the code is up to date with the latest MACE in this repo.
Using mace layers. This is a better option but still you need to implement the interactions loop, and it doesn't make it straightforward to reuse some pretrained weights from MACE.
Make MACE accept a readout function. I think this would not be as customizable and would complicate things inside MACE.
Just let it compute energies and not use them. Probably not too bad, but it would compute things unnecessarily.
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.
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 onMACE.forward
that defaults toTrue
could be great. Then one could set it toFalse
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