Batched evaluation with metatrain

Efficient evaluation of UPET models on a dataset is available from the command line via metatrain, which is installed as a dependency of UPET.

Step 1: export the model to TorchScript

Fetch and convert a UPET checkpoint from the HuggingFace repository:

mtt export https://huggingface.co/lab-cosmo/upet/resolve/main/models/pet-mad-s-v1.5.0.ckpt -o model.pt

Alternatively, fetch and save the model with the UPET Python API:

import upet

# Save the latest version of PET-MAD-S to a TorchScript file
upet.save_upet(
    model="pet-mad",
    size="s",
    version="1.5.0",
    output="model.pt",
)

Both commands download the model and convert it to TorchScript.

Step 2: write the evaluation options

Create an options.yaml file describing the dataset (in extxyz format) and the targets to predict:

systems: your-test-dataset.xyz
targets:
  energy:
    key: "energy"
    unit: "eV"

Step 3: run mtt eval

mtt eval model.pt options.yaml --batch-size=16 --output=predictions.xyz

This writes predictions.xyz with the predicted energies and forces for each structure in the dataset. For more options, see the metatrain evaluation docs.