# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
"""CLI entry point and programmatic API for NeuralBench.
Parse command-line arguments (``run_benchmark_cli``) or call the Python
API directly (``run_benchmark``). Task/model discovery, validation, and
config assembly are delegated to :mod:`neuralbench.registry` and
:mod:`neuralbench.experiment_config`.
"""
import argparse
import logging
import os
import sys
import traceback
import typing as tp
from neuralbench.experiment_config import build_experiment_configs
from neuralbench.plots._constants import DEFAULT_EVAL_MODE
from neuralbench.registry import (
ALL_DEVICES,
ALL_DOWNSTREAM_WRAPPERS,
ALL_MODELS,
ALL_TASKS,
ALL_UNVALIDATED_TASKS,
_format_datasets_epilog,
)
logger = logging.getLogger(__name__)
[docs]
def run_benchmark(
device: str,
task: str | list[str],
*,
model: str | list[str] | None = None,
dataset: str | list[str] | None = None,
space: str | list[str] | None = None,
checkpoint: str | None = None,
downstream_wrapper: str | list[str] | None = None,
grid: bool = False,
debug: bool = False,
max_neuro_hours_per_epoch: float | None = None,
force: bool = False,
retry: bool = False,
prepare: bool = False,
download: bool = False,
plot_cached: bool = False,
) -> list[dict[str, tp.Any]]:
"""Run one or more NeuralBench experiments from Python.
This is the programmatic equivalent of the ``neuralbench`` CLI: it
assembles experiment configs from the same YAML files and launches them.
For a model built outside this repo, and for the results of the runs in
hand, see :func:`neuralbench.evaluate_model`.
Parameters
----------
device : str
Brain recording device (``"eeg"``, ``"meg"``, ``"fmri"``, ...).
task : str or list of str
Task name(s), ``"all"``, or ``"all_multi_dataset"``.
model : str or list of str or None
Predefined model name(s), ``"all"``, ``"all_classic"``, ``"all_fm"``,
``"all_baseline"`` (chance / dummy / classical sklearn pipelines),
or ``None`` (uses default model from ``defaults/<device>/config.yaml``).
dataset : str or list of str or None
Dataset variant(s) or ``"all"``. ``None`` uses the base config.
space : str or list of str or None
Data-representation variant(s) from the task's ``spaces/`` folder, or
``"all"`` for the base config plus every variant. ``None`` uses each
model's ``default_space``, else the base config. Spaces outside a
dataset's ``available_spaces`` or a model's ``compatible_spaces`` are
skipped.
checkpoint : str or None
Path to a model checkpoint to reload.
downstream_wrapper : str or list of str or None
Adaptation-strategy preset name(s) from
``defaults/downstream_wrappers.yaml``, or ``"all"``. Swept over
foundation models only.
grid : bool
Expand the task-specific hyperparameter grid.
debug : bool
Run locally with a reduced config (2 epochs, 5 batches).
max_neuro_hours_per_epoch : float or None
Cap each training epoch to a fresh random draw of this many hours of
training windows, so large datasets train in comparable time.
force : bool
Force re-running experiments.
retry : bool
Retry failed experiments while keeping completed results.
prepare : bool
Run a single experiment to warm the preprocessing cache.
download : bool
Only download the dataset; do not run experiments.
plot_cached : bool
Generate plots and tables from cached results only, without
running any new experiments.
Returns
-------
list of dict
One result dict per experiment, with ``plot_cached=True`` only. Every
other mode launches experiments and returns an empty list, the results
being written to the results folder.
"""
logging.basicConfig(level=logging.INFO)
logging.getLogger("numexpr").setLevel(logging.WARNING)
# fontTools.subset emits very chatty INFO logs while matplotlib embeds
# font subsets into vector outputs (e.g. PDFs); mute them.
logging.getLogger("fontTools").setLevel(logging.WARNING)
logging.getLogger("fontTools.subset").setLevel(logging.WARNING)
logging.getLogger("fontTools.ttLib").setLevel(logging.WARNING)
logging.getLogger("matplotlib.font_manager").setLevel(logging.WARNING)
# exca and neuralset attach their own StreamHandlers; disable propagation
# to the root logger to avoid duplicate log lines.
logging.getLogger("exca").propagate = False
logging.getLogger("neuralset").propagate = False
if plot_cached and (force or retry or prepare):
raise ValueError(
"Cannot use force, retry, or prepare flags when plotting cached results."
)
configs = build_experiment_configs(
device,
task,
model=model,
dataset=dataset,
space=space,
checkpoint=checkpoint,
downstream_wrapper=downstream_wrapper,
grid=grid,
debug=debug,
max_neuro_hours_per_epoch=max_neuro_hours_per_epoch,
force=force,
retry=retry,
prepare=prepare,
download=download,
quiet=plot_cached,
)
if download:
return []
if plot_cached:
os.environ["CUDA_VISIBLE_DEVICES"] = ""
from neuralbench.main import BenchmarkAggregator
agg = BenchmarkAggregator(
# ConfDicts, which the pydantic model coerces into Experiment instances.
experiments=configs, # type: ignore[arg-type]
debug=debug,
device=device,
eval_mode=DEFAULT_EVAL_MODE if downstream_wrapper is None else None,
)
if not plot_cached:
agg.prepare()
results = []
if plot_cached:
logger.info("--- PREPARING GLOBAL PLOTS AND TABLES ---")
results = agg.run(cached_only=True)
return results
[docs]
def run_benchmark_cli() -> None:
"""CLI entry point for ``neuralbench``.
Parses command-line arguments and delegates to :func:`run_benchmark`.
"""
parser = argparse.ArgumentParser(
description="Run neuralbench.",
epilog=_format_datasets_epilog(),
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument(
"device",
type=str,
choices=ALL_DEVICES,
help="Brain device on which the desired task relies.",
)
parser.add_argument(
"task",
type=str,
nargs="+",
choices=["all", "all_multi_dataset"] + ALL_TASKS + ALL_UNVALIDATED_TASKS,
metavar="TASK",
help=(
"Task(s) to run. Use 'all' to run all tasks, "
"'all_multi_dataset' to run only tasks with multiple dataset variants, "
"or specify one or more task names (listed per device below)."
),
)
parser.add_argument(
"-g", "--grid", action="store_true", help="Run task-specific grid."
)
parser.add_argument(
"-d",
"--debug",
action="store_true",
help="Run in debug mode (locally and smaller config, with infra.mode='force').",
)
parser.add_argument(
"--max-neuro-hours-per-epoch",
type=float,
default=None,
metavar="HOURS",
help="Cap each training epoch to a fresh random draw of HOURS of training windows.",
)
mode_group = parser.add_mutually_exclusive_group()
mode_group.add_argument(
"-f", "--force", action="store_true", help="Force rerunning of experiment."
)
mode_group.add_argument(
"-r",
"--retry",
action="store_true",
help="Retry failed experiments (keep completed results).",
)
parser.add_argument(
"-p",
"--prepare",
action="store_true",
help="Run single experiment to prepare cache.",
)
model_choices = ["all", "all_classic", "all_fm", "all_baseline"] + ALL_MODELS
parser.add_argument(
"-m",
"--model",
nargs="*",
choices=model_choices,
metavar="MODEL",
help=(
"Override config to use one or more predefined models. Multiple models "
f"will be run in the grid. Choose from: {', '.join(model_choices)}."
),
)
parser.add_argument(
"-c",
"--checkpoint",
type=str,
help=(
"Path to a model checkpoint to reload. If this follows the format "
"`wandb:entity/project/grid`, it will be used to find all available checkpoint paths "
"for a specific wandb grid."
),
)
wrapper_choices = ["all"] + list(ALL_DOWNSTREAM_WRAPPERS.keys())
parser.add_argument(
"-w",
"--downstream-wrapper",
nargs="*",
choices=wrapper_choices,
metavar="WRAPPER",
help=(
"Adaptation strategy preset(s) to sweep over; applied to foundation "
f"models only. Choose from: {', '.join(wrapper_choices)}."
),
)
parser.add_argument(
"--download",
action="store_true",
help="Download the study. The experiment(s) will not be run.",
)
parser.add_argument(
"--pdb",
action="store_true",
help="Launch pdb on exception.",
)
parser.add_argument(
"--plot-cached",
action="store_true",
help="Plot from cached results only, without running any experiments.",
)
parser.add_argument(
"--dataset",
type=str,
default=None,
help=(
"Specify a dataset for the task, by study name (as listed below) or "
"by the stem of its file in the task's datasets/ folder, in any case. "
"Use 'all' to run on all available datasets. "
"A dataset other than the task's default merges datasets/<stem>.yaml "
"over the base config.yaml. "
"Example: --dataset Schalk2004Bci2000 or --dataset all"
),
)
parser.add_argument(
"--space",
type=str,
nargs="+",
default=None,
help=(
"Data representation(s) to run on, by the stem of their file in the "
"task's spaces/ folder (merged over the base config.yaml). "
"Use 'all' for the base config plus every space. Spaces a dataset "
"or model does not support are skipped with a warning. "
"Example: --space mni fsaverage5 or --space all"
),
)
args = parser.parse_args()
try:
run_benchmark(
device=args.device,
task=args.task,
model=args.model,
dataset=args.dataset,
space=args.space,
checkpoint=args.checkpoint,
downstream_wrapper=args.downstream_wrapper,
grid=args.grid,
debug=args.debug,
max_neuro_hours_per_epoch=args.max_neuro_hours_per_epoch,
force=args.force,
retry=args.retry,
prepare=args.prepare,
download=args.download,
plot_cached=args.plot_cached,
)
except Exception:
if not args.pdb:
raise
import pdb
tb = sys.exc_info()[2]
traceback.print_exc()
pdb.post_mortem(tb)
if __name__ == "__main__":
run_benchmark_cli()