[BE] Add smoke test for [escn,gemnet,equiformer_v2] train+predict, Add optimization test for [escn,gemnet,equiformer_v2]

[misko]

Add a test case for training a simple eSCN model on a small subset of OC22 (first 16 elements of train).

To suppor this I added the pickle dataset format. This is a simple list of pytorch geometric objects returned from the lmdb dataloaders.

def lmdb_dataset_to_pkl(lmdb_src, pickle_output_filename, n=128):
    ds = LmdbDataset(
        {
            "src": lmdb_src,
        }
    )
    pickle.dump([ds[idx] for idx in range(n)], open(pickle_output_filename, "wb"))

^ generate a pickle dataset from an existing lmdb dataset

We can use tests like this in the main repo and also in developing experimental branches to automatically test and confirm that we did not obviously break training.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 62.29%. Comparing base (daf72a5) to head (910aa86).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #640      +/-   ##
==========================================
+ Coverage   57.27%   62.29%   +5.01%     
==========================================
  Files         109      109              
  Lines       10316    10316              
==========================================
+ Hits         5909     6426     +517     
+ Misses       4407     3890     -517     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[mshuaibii]

Documenting our chat offline:

• Introducing another dataset is unnecessary, we can leverage existing readers even for the unit tests.
• Rather than uploading the datasets directly to the repo, we can upload them to s3 and download them as part of the tests. The same way we currently handle checkpoints.
• It would be great to include eqv2 instead or alongside escn as its the more commonly used model.
• We should make sure to test the other modes as well

[misko]

Documenting our chat offline:

• Introducing another dataset is unnecessary, we can leverage existing readers even for the unit tests.

I removed the pickle dataset, and used @r-barnes commit to use the tutorial dataset.

• Rather than uploading the datasets directly to the repo, we can upload them to s3 and download them as part of the tests. The same way we currently handle checkpoints.
• It would be great to include eqv2 instead or alongside escn as its the more commonly used model.

Added tests for eqv2 and gemnet

• We should make sure to test the other modes as well

Added test for predict mode (alongside the existing train mode test)

[lbluque]

Can you use a syrupy snapshot here instead of hardcoding the expected energy and force mae?

[misko]

Training is not reproducible to a reasonable accuracy :'( just due to numerical instability. In the above we want the metrics to be less than what is being specified not more. I dont think syrupy has a good way of doing this, but maybe i missed it

[lbluque]

Assuming this is a function so that it is easy to add additional train, val, test src to the tests.

Would it be better to make this a fixture (and if needed can it can be parametrized with different dataset src when the time comes)

[misko]

That was the intention, its only used in two ways in the tests and never both ways back to back.
One way is for training (train=tutorial_dataset,val=tutorial_dataset)
And the other in testing (train=tutorial_dataset,val=tutorial_dataset,test=tutorial_dataset)
I don't know how to parameterize it in some way to make it more readable. The examples i see online seem to show when we want to use all parameterizations of the fixture back to back.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

see 20 files with indirect coverage changes

[lbluque]

lgtm - thanks @misko!