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

[WIP] Introduce additional_labels #306

Open
wants to merge 83 commits into
base: main
Choose a base branch
from

Conversation

AntonioMirarchi
Copy link
Contributor

@AntonioMirarchi AntonioMirarchi commented Mar 15, 2024

This PR proposes several enhancements to enable the usage of extra arguments in both training and inference. The key modifications include

  • standardizing the naming conventions for all extra fields across different data loaders;
  • removing q and s as explicit args in the forward;
  • introducing additional labels through a new dictionary for model initialization, example:
additional_labels:
     tensornet_q:
        label: total_charge
        learnable: true
        init_value: 0.1
  • add TensornetQ nn.Module
  • extra args expansion from molecule-wise to atom-wise inside the model
  • tensornet_q tested with total_charge and partial_charges, reproducibility achieved
  • rewriting some tests according to the new methods
  • allowing multiple additional_methods implementation
    At the moment only 'tensornet_q' supported as additional_methods but more can be added

torchmdnet/models/tensornet.py Outdated Show resolved Hide resolved
torchmdnet/models/tensornet.py Outdated Show resolved Hide resolved
torchmdnet/models/tensornet.py Outdated Show resolved Hide resolved
torchmdnet/models/tensornet.py Outdated Show resolved Hide resolved
torchmdnet/models/tensornet.py Outdated Show resolved Hide resolved
torchmdnet/models/torchmd_gn.py Show resolved Hide resolved
torchmdnet/models/torchmd_t.py Show resolved Hide resolved
torchmdnet/models/tensornet.py Outdated Show resolved Hide resolved
torchmdnet/module.py Outdated Show resolved Hide resolved
torchmdnet/scripts/train.py Show resolved Hide resolved
@RaulPPelaez
Copy link
Collaborator

It would be wise to add a new test file in which you poke the new functionality around, for instance making sure that the input is being received correctly or that the output of the model changes depending on whether additional_labels is passed or not.

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.

None yet

3 participants