Restart a training changing ForceOutput with StressForceOutput

[LucaBrugnoli]

I trained a model on forces and energies and then I used it to perform dynamics simulations with LAMMPS. While NVT dynamics seem to be stable, NPT dynamics are unstable and explode. Following this discussion , I tried to use the stress branch of pair_nequip with the same model, but as expected, I encountered a compatibility error due to the lack of stress encoding in the model.
Then, I attempted to 'update' my model with stress with a nequip_training job by initializing the previous model state with initial_model_stateand initialize_from_state and replacing ForceOutput with StressForceOutput in model_builders, without making any other changes to the YAML config. However, the job fails after reloading the dataset, just before generating the network.
I would like to ask whether this is due to an incompatibility with the type of initialization I want to perform or if I am missing something in the YAML config.

[Linux-cpp-lisp]

Hi @LucaBrugnoli ,

Yes, this sounds like the correct process, although you may need to use the load_model_state model builder and the load_model_state: /path/to/somewhere.pth option, not initialize_from_state.

However, the job fails after reloading the dataset,

Can you please provide the exact, full error output?

[LucaBrugnoli]

Here it is the error message, after the loading of the dataset.

Successfully loaded the data set of type ASEDataset(41119)...
Replace string dataset_forces_rms to 37.753509521484375
Replace string dataset_per_atom_total_energy_mean to -7551.79931640625
Atomic outputs are scaled by: [O, Si: 37.753510], shifted by [O, Si: -7551.799316].
/ccc/work/cont003/gen7069/brugnolu/My_venv/lib/python3.8/site-packages/nequip/nn/_grad_output.py:196: UserWarning: !! Stresses in NequIP are in BETA and UNDER DEVELOPMENT: please carefully check the sanity of your results and report any (potential) issues on the GitHub
warnings.warn(
Replace string dataset_forces_rms to 37.753509521484375
Initially outputs are globally scaled by: 37.753509521484375, total_energy are globally shifted by None.
Traceback (most recent call last):
File "/ccc/work/cont003/gen7069/brugnolu/My_venv/bin/nequip-train", line 8, in
sys.exit(main())
File "/ccc/work/cont003/gen7069/brugnolu/My_venv/lib/python3.8/site-packages/nequip/scripts/train.py", line 72, in main
trainer = fresh_start(config)
File "/ccc/work/cont003/gen7069/brugnolu/My_venv/lib/python3.8/site-packages/nequip/scripts/train.py", line 159, in fresh_start
final_model = model_from_config(
File "/ccc/work/cont003/gen7069/brugnolu/My_venv/lib/python3.8/site-packages/nequip/model/_build.py", line 89, in model_from_config
model = builder(**params)
File "/ccc/work/cont003/gen7069/brugnolu/My_venv/lib/python3.8/site-packages/nequip/model/_weight_init.py", line 26, in initialize_from_state
return load_model_state(
File "/ccc/work/cont003/gen7069/brugnolu/My_venv/lib/python3.8/site-packages/nequip/model/_weight_init.py", line 52, in load_model_state
model.load_state_dict(state, strict=config.get(_prefix + "_strict", True))
File "/ccc/work/cont003/gen7069/brugnolu/My_venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1604, in load_state_dict
raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for RescaleOutput:
Missing key(s) in state_dict: "model._empty".
srun: error: irene7051: task 0: Exited with exit code 1

• exit 0

I will give a try with load_model_state

[Linux-cpp-lisp]

To avoid this error set load_model_state_strict: false when using load_model_state.

[LucaBrugnoli]

It worked fine! Thanks!