search

lab-cosmo / metatrain

Training and evaluating machine learning models for atomistic systems.
https://lab-cosmo.github.io/metatrain/
BSD 3-Clause "New" or "Revised" License
13 stars 3 forks source link

Remove excessive branching from `eval_model` logic #249

Closed frostedoyster closed 3 weeks ago

frostedoyster commented 3 weeks ago

Closes #215

📚 Documentation preview 📚: https://metatrain--249.org.readthedocs.build/en/249/

frostedoyster commented 3 weeks ago

I tried to restructure the code, but I get this:

python3.12/site-packages/metatrain/utils/evaluate_model.py:300: UserWarning: This system's positions or cell requires grad, but the neighbors does not. You should use `register_autograd_neighbors()` to make sure the neighbors distance vectors are integrated in the computational graph. (Triggered internally at /home/filippo/code/demo/metatensor-models/metatensor/metatensor-torch/src/atomistic/system.cpp:616.)

Should I hide it?

Luthaf commented 3 weeks ago

From looking around the code, _eval_targets in eval.py is now computing the NL twice: once line 160 of eval.py, and one in the changes of this PR. I would remove the one in eval.py!

frostedoyster commented 3 weeks ago

I don't think we're recomputing the NL in the changes that are introduced in this PR. We're just re-building the System object (taking all NLs that were attached to it and attaching them to the new object). No NL calculation should be done here

frostedoyster commented 3 weeks ago

Unfortunately I get this warning now on all architectures for evaluation:

/home/filippo/code/metatrain-test/virtualenv/lib/python3.12/site-packages/metatrain/utils/evaluate_model.py:300: UserWarning: This system's positions or cell requires grad, but the neighbors does not. You should use `register_autograd_neighbors()` to make sure the neighbors distance vectors are integrated in the computational graph. (Triggered internally at /project/metatensor-torch/src/atomistic/system.cpp:615.)
  new_system.add_neighbor_list(nl_options, nl)

Moreover, PET (which is the only architecture that uses the precomputed interatomic vectors from the NL) crashes, probably due to the warning above. PET runs fine on main however

frostedoyster commented 3 weeks ago

@Luthaf it works on all architectures now, thank you!