Benchmark Several Models#
The xnn.common.benchmark package scores any set of pre-trained
models on one dataset and writes a table of error metrics. It only does one
thing: benchmarking. Thus, no other operation such as data processing and model
training is performed in this module. The users must produce the model
checkpoints first (for example with xnn train; see
Train a Model), and then include the checkpoint location in the
config. Then the benchmark module will load the model from the checkpoint, run
it on the dataset, and report the metrics.
Write a benchmark config#
A benchmark config lists the models to compare, the data to score on,
and the metrics to report:
# benchmark.yaml
models:
- label: mace
checkpoint: runs/mace/best.pt # architecture read from the checkpoint
- label: nequip
checkpoint: runs/nequip/best.pt
metrics: # target -> its reported metrics
energy: [mae, rmse] # energy / forces / stress
forces: [mae, rmse] # mae / mse / rmse or custom
data:
test_path: data/argon_test.extxyz # the dataset to score on
batch_size: 16
output:
dir: runs/benchmark
filename: results
formats: [csv, json, md] # any registered writer
device: auto
seed: 1234
The bundled configs/benchmark.yaml is a complete, commented example, and
examples/benchmark/ benchmarks NequIP, Allegro, MACE and PhysNet on the
Argon test set.
Every model entry needs a checkpoint: The xnn.common.benchmark
module scores pre-trained models. An xnn checkpoint embeds the
Config it was trained with, so the benchmark
rebuilds the exact architecture from the checkpoint. A model entry usually needs
only its checkpoint, and an optional label for the row name. The users
should provide an explicit architecture – inline keys, a name, or a
config file – only for checkpoints that do not embed a config object.
The upstream key spellings (e.g., for MACE, r_max, …) are translated to
the xnn canonical names exactly as in a normal run (see
Reuse Upstream MACE/NequIP Configs).
The dataset is resolved from the data section: test_path (the natural
held-out benchmark set), falling back to val_path then train_path.
Connect metrics to targets#
The metrics mapping states unambiguously which metric is reported for
which target quantity, so different targets can carry different metrics.
Targets are any subset of energy / forces / stress (anything else is
rejected at load time). Only the targets named in the mapping are scored. Three
equivalent spellings are accepted. The mapping is the canonical form, in which
the targets are keys and the metrics, corresponding to each target, can be a
single name or a list of names:
metrics:
energy: mae
forces: [mae, rmse]
The same connection can be written as (target, metrics) pairs: 2-item lists
in YAML are turned into tuples in Python. Repeated targets accumulate their
metrics:
metrics:
- [energy, mae]
- [forces, [mae, rmse]]
Finally, the flat cross-product shorthand applies every metric to every target:
metrics: [mae, rmse]
targets: [energy, forces]
The separate targets key belongs to the flat shorthand only. combining it
with the mapping or pair form raises an error. An omitted side of the flat form
falls back to its default (metrics: [mae, rmse], targets: [energy,
forces]), as does a mapping target given without metrics (e.g. energy:
null). Whatever the spelling, every loader normalizes the inputs into the same
per-target mapping on BenchmarkConfig.metrics, and the results table carries one
<target>_<metric> column per connected pair.
Run the benchmark#
From the command line:
xnn benchmark --config benchmark.yaml
xnn benchmark --config benchmark.yaml --set "metrics={'energy': ['mae']}"
or from Python:
from xnn.common.benchmark import from_yaml, run_benchmark
rows = run_benchmark(from_yaml("benchmark.yaml"))
The comparison table is printed and written to output.dir in every
configured format, one row per model, with a n_params column showing the
number of model parameters and one <target>_<metric> column per scored
quantity (e.g. energy_mae, forces_rmse).
The tables show the physical unit after each metric name (e.g. energy_mae
[eV/atom]). xnn is unit-agnostic, so these are labels only: they default to
eV/atom (or eV with energy_per_atom: false) for energy, eV/A for
forces and eV/A**3 for stress, and are overridden per target with units
to match your data:
units:
energy: meV/atom
forces: meV/A
The written results.{csv,json,md} files carry the same unit-annotated
headers (the returned rows keep plain keys for programmatic use).
Report atomization (interaction) energy#
By default energy is scored per atom (total energy divided by the atom count).
Set atomic_energies to score the atomization / interaction energy
instead: the total energy minus the summed per-element reference energies (E0s),
which is the physically meaningful quantity:
atomic_energies: {1: -13.663, 6: -1029.863, 8: -2042.785} # {Z: E0}, eV
# atomic_energies: average # or fit E0s from the benchmark data (lstsq)
# species: [H, C, O] # only needed for the list / scalar forms
energy_per_atom: true # divide the (atomization) energy by atom count
The same E0 offset is subtracted from both the prediction and the reference, so
plain difference metrics (MAE/MSE/RMSE) are unchanged while the reported values
become meaningful; the effect is visible in reference-dependent custom metrics
(relative error, \(R^2\)) and in reported magnitudes. Accepted forms are a
{Z: E0} / {symbol: E0} dict, a list aligned with species, a single
number, or "average" to fit the E0s from the benchmark dataset
(see build_e0_lookup() for details ).
Add custom metrics and output formats#
Custom error metrics are registered like models. Either import them in the
config with a module:function path (each callable has the signature
fn(pred, target) -> float):
custom_metrics:
# import a metric from a Python module
- {name: maxae, path: my_metrics:max_abs_error}
metrics:
energy: [mae, rmse]
# use the custom metric by name
forces: [mae, rmse, maxae]
or register them in Python before running:
from xnn.common.benchmark import register_metric, register_writer
@register_metric("maxae")
def max_abs_error(pred, target):
return float((pred - target).abs().max())
@register_writer("tsv") # a new output format
def write_tsv(rows, columns, path):
with open(path, "w") as f:
f.write("\t".join(columns) + "\n")
for r in rows:
f.write("\t".join(str(r.get(c, "")) for c in columns) + "\n")
Any registered writer name can then be listed under output.formats.