xnn.common.benchmark.runner#

Score pre-trained models on one dataset and tabulate their errors.

Benchmarking does one thing: for every model in the config it builds the architecture, loads the entry’s checkpoint into it, scores it on the benchmark dataset with the configured metrics, and writes a comparison table. It does not train or evaluate during training – produce the checkpoints with xnn train (or any other route) first.

Model building reuses the model registry and ForceStressOutput, so a benchmarked model is built and run exactly as it would be in a single run. The benchmark dataset is built once per cutoff, so models sharing a cutoff share one neighbor-list build.

Two independent kinds of parallelism are available, and they combine:

  • Across the dataset (one model, several GPUs or nodes). Launched through a distributed launcher exactly like training (torchrun --nproc-per-node 4 -m xnn benchmark ...), every rank scores a disjoint strided shard of the dataset and the predictions are gathered before the metrics are computed, so the numbers equal those of a single-process run. Only rank 0 writes.

  • Across models (one GPU per model). select_entries() restricts a run to some entries (--models / --shard), each such run writes its rows to a part file, and merge_parts() assembles the final table from the parts. On a cluster the parts are the tasks of a Slurm job array; on one node run_parallel() spawns one worker process per GPU and merges itself.

Functions

column_units(cfg, rows)

Map each metric column of rows to its unit label (see Benchmark._column_units()).

merge_parts(cfg)

Assemble the final table from the part files of a split benchmark.

part_tag(entries)

A filesystem-safe tag naming a subset of entries (their labels joined by +).

run_benchmark(cfg)

Convenience wrapper: build a Benchmark and run it.

run_parallel(raw[, workers])

Score the models of a benchmark concurrently, one worker process per model.

select_entries(cfg[, models, shard])

Pick the model entries one run of a split benchmark scores.

write_part(cfg, rows, tag)

Write the rows of one run of a split benchmark to its part file.

write_report(cfg, rows)

Print the results table and write it to every configured format.

Classes

Benchmark(cfg[, entries, part])

Run a BenchmarkConfig and collect a comparison table.
Show Source