Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement ZBL potential #134

Merged
merged 13 commits into from
Nov 9, 2022
Merged

Implement ZBL potential #134

merged 13 commits into from
Nov 9, 2022

Conversation

peastman
Copy link
Collaborator

@peastman peastman commented Oct 6, 2022

This is the first piece of #26. It isn't fully tested yet, but I think it's ready for a first round of comments.

I created a mechanism for passing arguments to the prior. prior_args is now the option specified by the user in the config file. prior_init_args stores the value returned by get_init_args(), which contains the arguments needed for reconstructing it from a checkpoint.

This prior requires the dataset to provide several pieces of information. To keep things general, the HDF5 format allows the file to contain a _metadata group which can store arbitrary pieces of information. Most of the other dataset classes should be able to hardcode the necessary values, since they aren't intended to be general.

@stefdoerr
Copy link
Collaborator

Looks good to me. Tests are failing though:

 E       RuntimeError: 
E       'NoneType' object has no attribute or method 'atomwise'.:
E         File "/usr/share/miniconda/envs/torchmd-net/lib/python3.9/site-packages/torchmdnet/models/model.py", line 182
E           
E               # apply atomwise prior model
E               if self.prior_model is not None and self.prior_model.atomwise:
E                                                   ~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
E                   x = self.prior_model(x, z, pos, batch)

Also you might have some merge conflicts with the latest PR of Raimondas

@stefdoerr
Copy link
Collaborator

stefdoerr commented Oct 7, 2022

Wait a second, how does this error even occur? Shouldn't the if short-circuit if self.prior_model is not None? Weird. Maybe jit-ed conditionals don't short-circuit?

torchmdnet/models/model.py Outdated Show resolved Hide resolved
torchmdnet/models/model.py Outdated Show resolved Hide resolved
torchmdnet/priors.py Outdated Show resolved Hide resolved
@giadefa
Copy link
Contributor

giadefa commented Oct 7, 2022 via email

@peastman
Copy link
Collaborator Author

peastman commented Oct 7, 2022

Who runs the model knows what units want to use and set it for the priors

The units are determined by the dataset. It isn't a free choice for the user to make. If your dataset contains positions in nm and energies in kJ/mol, that's what the prior needs to work with. Any other units would produce wrong results.

We could consider creating an automated unit conversion system like I suggested in #26 (comment), but that's a separate project.

@peastman
Copy link
Collaborator Author

Wait a second, how does this error even occur? Shouldn't the if short-circuit if self.prior_model is not None?

Yes it should! I think it's a bug in the torchscript jit. I was able to work around it by splitting it into two lines.

@peastman peastman changed the title [WIP] Implement ZBL potential Implement ZBL potential Oct 29, 2022
@peastman
Copy link
Collaborator Author

This is ready for review.

@stefdoerr
Copy link
Collaborator

The test is failing if you could take a look

torchmdnet/scripts/train.py Outdated Show resolved Hide resolved
@peastman peastman mentioned this pull request Oct 31, 2022
3 tasks
@peastman
Copy link
Collaborator Author

The test is failing if you could take a look

Fixed. I couldn't reproduce it locally, but I managed to guess what it was unhappy about.

Copy link
Collaborator

@stefdoerr stefdoerr left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fine by me. Does anyone else want to comment?

@peastman peastman mentioned this pull request Nov 1, 2022
raimis
raimis previously requested changes Nov 2, 2022
torchmdnet/priors/zbl.py Outdated Show resolved Hide resolved
@@ -43,7 +43,7 @@ def test_train(model_name, use_atomref, tmpdir):
prior = None
if use_atomref:
prior = getattr(priors, args["prior_model"])(dataset=datamodule.dataset)
args["prior_args"] = prior.get_init_args()
args["prior_init_args"] = prior.get_init_args()
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The loading of the pretrained models (https://github.com/torchmd/torchmd-net/tree/main/examples#loading-checkpoints) fails:

from torchmdnet.models.model import load_model
load_model('ANI1-equivariant_transformer/epoch=359-val_loss=0.0004-test_loss=0.0120.ckpt')

AssertionError: Requested prior model Atomref but the arguments are lacking the key "prior_init_args".

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It sounds like we want to redo how prior args are specified as described in #26 (comment). That means we can switch back to prior_args for this value. There will still be compatibility issues, because it will need to become a list with args for multiple prior models, but I can add a check for that case for backward compatibility. I'll go ahead and make the changes in this PR.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Might be a good idea to add a test case for loading model checkpoints from a previous version.

@peastman
Copy link
Collaborator Author

peastman commented Nov 3, 2022

I added the multiple priors support. I tried to make the syntax fairly flexible. All of the following are valid.

prior_model: Atomref
prior_model:
  - ZBL:
      cutoff_distance: 4.0
      max_num_neighbors: 50
  - Atomref
prior_model:
  - ZBL:
      cutoff_distance: 4.0
      max_num_neighbors: 50
    Atomref:

I added tests for creating models both from config files and from checkpoints. And the test for older checkpoints in test_module.py continues to work. I can't be sure I've caught all possible problems though, so if you can check your own files that would be helpful.

@peastman peastman dismissed raimis’s stale review November 8, 2022 17:28

The mechanism has been replaced with a different one.

@peastman
Copy link
Collaborator Author

peastman commented Nov 8, 2022

Is this ok to merge?

Copy link
Collaborator

@raimis raimis left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is good to go!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants