Training stress with allegro

[Tawfiq1448]

Hello,

I am trying to train Allegro model with stress. I have made the following modifications in my config file:

1. I changed the 'ForceOutput' in the model_builders to 'StressForceOutput'
2. I also included the virial loss_coeff as: virial: 0.8

I tried to train this separately with NVT and NPT data. Both times I encountered an error failing the equivariance test. The error is as follows:

AssertionError: Equivariance test failed for cases:
   (parity_k=1, did_translate=False, field=stress              ) -> max error=2.352e-03
   (parity_k=1, did_translate=True , field=stress              ) -> max error=2.352e-03

I am using nequip==0.5.5, Torch==1.11.0+cu113, and mir-allegro==0.2.0. Can you please tell me a solution for this?

I have another question regarding hyperparameter tuning. I am trying to find good hyperparameter combinations of number of layers, batch number, learning rate, cutoff, and number of features for my dataset. For this purpose, if I use the nequip-benchmark, it will help me to understand the model's performance for my future MD calculations. If I aim to understand the effect of hyperparameters on model errors (validation and training error) and thereby get an optimized hyperparameter combination for the lowest error, do I need to use nequip-train? If I need to use nequip-train, can I do 10 epochs for all hyperparameter combinations to compare the MAE and RMSE at 10th epoch for all combinations of hyperparameters? I am asking this because it will save a good amount of computational time as I will not need to do the full training until convergence for all these hyperparameter combinations.

[tianyus5]

I had the same problem...

[Tawfiq1448]

Update: I solved the first part of the problem here regarding training with stress. It looks like stress does not pass equivariance test. So I removed the equivariance-test tag from the nequip-train command.
Still, I would love to get any input you can provide regarding hyperparameter tuning. Thanks.

[Linux-cpp-lisp]

Can you send the full output of the equivariance test?

[Tawfiq1448]

Not sure if you are asking for this. This is my output file.

output.log

[Linux-cpp-lisp]

Hm... which version of nequip are you using, and if it's not the develop branch, can you try the develop branch?

[Tawfiq1448]

Sorry for the delayed response as I needed time to get back on this. With the develop branch of nequip, stress passes the equivariance test. Thanks.

[Linux-cpp-lisp]

Great, thank you for letting us know that this is resolved by develop!

[Linux-cpp-lisp]

If I need to use nequip-train, can I do 10 epochs for all hyperparameter combinations to compare the MAE and RMSE at 10th epoch for all combinations of hyperparameters?

My understanding is that this is a bit of an open question. Certainly you don't need to converge models all the way to get a sense of the relative ranking of hyperparameters, but how many epochs that takes will depend on the system and the size of your dataset (one epoch of training is very different in a 1M frame dataset and a 100 frame one).

[Tawfiq1448]

Thanks for your reply. This resonates with my thoughts too.