[WIP] Introduce additional_labels - Githubissues

torchmd / torchmd-net

Training neural network potentials

MIT License

326 stars 73 forks source link

[WIP] Introduce additional_labels #306

Open AntonioMirarchi opened 6 months ago

AntonioMirarchi commented 6 months ago

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

[x] standardizing the naming conventions for all extra fields across different data loaders;
[x] removing q and s as explicit args in the forward;

[x] introducing additional labels through a new dictionary for model initialization, example:

additional_labels:
 tensornet_q:
    label: total_charge
    learnable: true
    init_value: 0.1

[x] add TensornetQ nn.Module
[x] extra args expansion from molecule-wise to atom-wise inside the model
[x] tensornet_q tested with total_charge and partial_charges, reproducibility achieved
[x] rewriting some tests according to the new methods
[x] allowing multiple additional_methods implementation At the moment only 'tensornet_q' supported as additional_methods but more can be added

RaulPPelaez commented 6 months ago

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.