search

FAIR-Chem / fairchem

FAIR Chemistry's library of machine learning methods for chemistry
https://opencatalystproject.org/
Other
773 stars 242 forks source link

How to add a new model to fair-chem #732

Closed MSJavaScript closed 2 months ago

MSJavaScript commented 3 months ago

I want to add my own nn-model to the fairchem, with my own collate function and loss function, how can I do that ?

MSJavaScript commented 3 months ago

ps: I want to train the model on multi-core CPU or a single GPU

mshuaibii commented 3 months ago

Hi -

Adding and using your own model is as simple as adding it here - https://github.com/FAIR-Chem/fairchem/tree/main/src/fairchem/core/models and ensuring it's added to the registry by including the following - https://github.com/FAIR-Chem/fairchem/blob/main/src/fairchem/core/models/schnet.py#L19.

As far as loss function, if you are using a built-in pytorch loss function you just need to specify it in the config - https://github.com/FAIR-Chem/fairchem/blob/657598b4ed301b589601016397e8f4012db7d3dd/configs/ocp_example.yml#L137. Otherwise you'll need to manually add it here - https://github.com/FAIR-Chem/fairchem/blob/657598b4ed301b589601016397e8f4012db7d3dd/src/fairchem/core/trainers/base_trainer.py#L526-L545