Load a Pre-trained Model#
The from_pretrained() function loads the model artifacts: the original
foundation models pre-trained by their main developer teams (e.g., MACE and
AIMNet2), models trained with xnn, Zenodo uploads and local checkpoints, all can
be loaded through the same function call and the same cache. The
examples/models/pretrained_models_tutorial.ipynb notebook provides a
runnable version of this page.
Inspect the model registry#
from xnn.common.models import list_models, model_card
# List all registered and cached models
list_models()
# ['mace-off23-large', 'mace-off23-medium', 'mace-off23-small']
list_models(tag="organic")
# Source, license, citation, cutoff, elements, heads
model_card("mace-off23-small")
The xnn models command does the same from a shell:
# Models that load without a download
xnn models list --cached
# The model card
xnn models info mace-mh-0
Load a model#
from xnn.common.models import from_pretrained
model = from_pretrained("mace-off23-small") # Registered name
model = from_pretrained("mace-mh-0", head="omat_pbe") # One head of a multi-head model
model = from_pretrained("aimnet2") # AIMNet2 with its D3 term
model = from_pretrained("doi:10.5281/zenodo.18957344", # Zenodo DOI
filename="mace_csfapbbri_al_5_1_stagetwo.model")
model = from_pretrained("runs/exp/best.pt") # Trainer checkpoint
model = from_pretrained("path/to/model_dir") # Portable model directory
The from_pretrained function returns a model object wrapped in the
ForceStressOutput class. This object is
already in the eval() mode and in the model checkpoint’s dtype.
Here are some additional options for customizing the load behavior:
from_pretrained(name, wrap=False) # The bare potential
from_pretrained(name, dtype=torch.float64) # Cast the weights
from_pretrained(name, device="cuda") # Move it
from_pretrained(name, compute_stress=True) # Add the stress head
from_pretrained(name, dispersion="d4") # Add D3 / D4 to a plain model
from_pretrained(name, use_fast=True) # Fused GPU kernels
The load_pretrained() function behaves similar to
from_pretrained() but returns a
PretrainedModel which also provides its
config, model card and other information. For ASE, you can wrap the model in an
XNNCalculator using:
from xnn.common.deploy import XNNCalculator
atoms.calc = XNNCalculator.from_pretrained("mace-mp-0-medium")
The cache#
Each model is cached once into a portable directory with card.json,
config.yaml and model.pt (no pickled objects, no absolute paths). An
external format such as a mace-torch file is automatically converted to an
xnn-native format on the first load and subsequent loads will be carried out
offline.
from_pretrained("mace-off23-small", cache_dir="/scratch/me/xnn-models")
Without cache_dir the XNN_MODELS environment variable sets the
directory, then $XNN_CACHE/models, and then models/ in a source checkout
(~/.cache/xnn/models for an installed package).
The XNN_OFFLINE=1 or local_files_only=True variables forbid xnn to
access the network. Setting force_download=True forces xnn to fetch the
model over the network again.
Cached directories are designed to be portable and self-sufficient. You can copy them to another machine and they can be loaded there by using the folder name, in offline mode.