Equitrain: Training and Fine-Tuning Machine Learning Interatomic Potentials

Equitrain is a Python toolkit for preprocessing atomistic datasets, training machine-learning interatomic potentials (MLIPs), fine-tuning existing checkpoints, and running evaluation or prediction through one CLI/API.

Features

• Unified Torch and JAX training entry points.
• Model wrappers for MACE, SevenNet, ORB, ANI, and M3GNet.
• Native HDF5 preprocessing for large atomistic datasets.
• Torch reaction-relative losses for barrier and reaction energies.
• Fine-tuning adapters for Delta/L²-SP, Freeze, and LoRA workflows.
• ASE calculator helpers for batched prediction and relaxation.

Supported Models

For MACE, use mace-model for model construction/conversion and equitrain for preprocessing, training, fine-tuning, checkpointing, evaluation, and prediction.

Documentation

Full documentation is published at https://bamescience.github.io/equitrain/:

• Installation
• Quickstart
• Data and Preprocessing
• CLI
• Training Options
• Python API
• API Reference
• Model Wrappers
• JAX Bundles
• Fine-Tuning
• Phonon Fine-Tuning Paper
• Calculators
• Reaction-Relative Losses
• Resources

The documentation source is in docs/. Build or serve it locally with:

pip install -e '.[docu]'
mkdocs serve

Optional EquiLoad HDF5 Loader

Equitrain's portable PyTorch/PyG loader remains the default. EquiLoad is a separately installed native extension. This integration is tested with the v0.1.0 release. Build it against the exact Torch installation in the Equitrain environment, because PyTorch's C++ ABI is not stable between Torch releases:

git clone --branch v0.1.0 https://github.com/BAMeScience/EquiLoad.git equiload
cd equiload
bash scripts/install-python.sh --venv /path/to/equitrain/.venv
cd ..
equitrain --backend torch --torch-loader equiload ...

Use the tagged release when building EquiLoad. Earlier development checkouts also reported version 0.1.0 but lacked the global sample-index fix required for correct multi-file loading. Rebuild those older installations from the tag.

Do not install a wheel built in another environment. EquiLoad is CPU-side; Equitrain moves its returned batches to the active accelerator device. It currently requires Torch 2.6 through 2.14; the standard Python loader remains available for other supported Equitrain Torch versions.

The EquiLoad loader is currently single-process and does not support weighted resampling or relative reaction losses. It returns the same nested PyG batch structure as the standard loader and moves batches to Equitrain's active accelerator device. Without equiload installed, the default --torch-loader python requires no Rust tooling.

Installation

pip install equitrain

Until the package is fully available on PyPI, install from a local clone:

git clone https://github.com/BAMeScience/equitrain.git
cd equitrain
python3.10 -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
pip install uv
uv pip install -e '.[dev,docu]'

Install model/runtime extras as needed:

pip install 'equitrain[torch,mace]'
pip install 'equitrain[jax,mace-jax]'
pip install 'equitrain[torch,ani]'

Minimal Workflow

Preprocess data:

equitrain-preprocess \
    --train-file data-train.xyz \
    --valid-file data-valid.xyz \
    --compute-statistics \
    --atomic-energies average \
    --output-dir data \
    --r-max 4.5

Train a Torch/MACE model:

equitrain -v \
    --train-file data/train.h5 \
    --valid-file data/valid.h5 \
    --output-dir runs/mace \
    --model path/to/mace.model \
    --model-wrapper mace \
    --epochs 10 \
    --tqdm

Evaluate and predict:

equitrain-evaluate -v \
    --test-file data/test.h5 \
    --model path/to/mace.model \
    --model-wrapper mace \
    --output-dir evaluation_mace

equitrain-predict \
    --predict-file data/valid.h5 \
    --model path/to/mace.model \
    --model-wrapper mace \
    --output-dir predictions_mace

See the Quickstart for the full workflow, including JAX bundles and fine-tuned checkpoint export.

Fine-Tuning Note

Equitrain's Delta adapter is a residual-parameter implementation of L²-SP ("Starting Point") regularization from Li, Grandvalet, and Davoine, 2018, Explicit Inductive Bias for Transfer Learning with Convolutional Networks. It parameterizes fine-tuning as theta = theta_0 + delta, so weight decay on trainable deltas regularizes ||delta||_2^2.

Delta combined with freeze_layers is targeted L²-SP (L²-TSP): the L²-SP penalty applies only to selected trainable delta layers while frozen layers remain exactly at their pre-trained starting values. See Fine-Tuning.

Resources

Example data-preparation scripts are in resources/data, training scripts are in resources/training, and initial model examples are in resources/models.