Error training on OMol25

[peastman]

I'm trying to use the finetune.py script for training a model with OMol25. As an initial proof of concept, I downloaded the 4M subset from https://huggingface.co/facebook/OMol25/blob/main/DATASET.md. I then converted just a single one of the database files to SQLite format with

ase convert data0000.aselmdb data.db

I then tried to run the script with

python finetune.py --data_path=data.db --base_model=orb_v3_conservative_omol

It fails with an exception:

Original Traceback (most recent call last):
  File "/home/peastman/miniconda3/envs/mace/lib/python3.12/site-packages/torch/utils/data/_utils/worker.py", line 349, in _worker_loop
    data = fetcher.fetch(index)  # type: ignore[possibly-undefined]
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/peastman/miniconda3/envs/mace/lib/python3.12/site-packages/torch/utils/data/_utils/fetch.py", line 50, in fetch
    data = self.dataset.__getitems__(possibly_batched_index)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/peastman/miniconda3/envs/mace/lib/python3.12/site-packages/orb_models/common/dataset/abstract_dataset.py", line 65, in __getitems__
    return [self[idx] for idx in indices]
            ~~~~^^^^^
  File "/home/peastman/miniconda3/envs/mace/lib/python3.12/site-packages/orb_models/common/dataset/ase_sqlite_dataset.py", line 82, in __getitem__
    atoms.info.update(self.target_config.extract(row, self.name, "targets"))
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/peastman/miniconda3/envs/mace/lib/python3.12/site-packages/orb_models/common/dataset/property_definitions.py", line 346, in extract
    name: p.row_to_property_fn(row=row, dataset=dataset_name)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/peastman/miniconda3/envs/mace/lib/python3.12/site-packages/orb_models/common/dataset/property_definitions.py", line 175, in stress_row_fn
    raise ValueError(
ValueError: db row 53824 of mp-traj doesn't have an stress attribute and we haven't defined an alternative method to extract stress info.

I tried also adding the option --stress_loss_weight=0, but it didn't help.

What do I need to do differently? And is there documentation on how the OMol based models were trained so I can reproduce it?

[vsimkus]

Thanks for flagging this! This is related to a recent PR that enabled stress computation for omol convervative models. I've made a fix in #155, so setting --stress_loss_weight=0 will work!

[peastman]

Thanks!

Is there any documentation on how the models were trained so I can reproduce the protocol that was used?

[timduignan]

Thanks for catching that! Everything we've released is in this paper: https://arxiv.org/pdf/2504.06231 and this blog post: https://www.orbitalindustries.com/posts/orbmol-extending-orb-to-molecular-systems

[peastman]

I looked at both of those, but neither includes a lot of the basic details needed to reproduce it: the optimizer used, the learning rate schedule, the weights of the terms in the loss function, the number of epochs, etc.

[benrhodes26]

I am afraid we don't release the original training code or hyperparameters.