update benchmarks CLI

bash/run_docker.sh

@@ -19,7 +19,7 @@ DATASET="/workdir/dataset/"
 # -v : mount a volume from host machine at $MOUNT_POINT to container's /root/mne_data/
 # "${TAG}" : use the image with the tag 'moabb'
 # /usr/bin/python : use python command
-# /workdir/moabb/run.py : run the script 'run.py' located at '/workdir/moabb/'
+# /workdir/moabb/benchmark.py : run the script 'benchmark.py' located at '/workdir/moabb/'
 # --pipeline $PIPELINE : use pipelines located at $PIPELINE
 # --results $RESULTS : store results in $RESULTS
 # --output $OUTPUTS : store outputs in $OUTPUTS
@@ -30,7 +30,7 @@ docker run -it \
     -v "${MOUNT_POINT}/output/:/workdir/output" \
     "$TAG" \
     /usr/bin/python \
-    /workdir/moabb/run.py \
+    /workdir/moabb/benchmark.py \
     --pipeline "$PIPELINE" \
     --results "$RESULTS" \
     --output "$OUTPUTS" \

moabb/benchmark.py

@@ -2,13 +2,15 @@
 import os
 import os.path as osp
 from pathlib import Path
+from typing import List, Union
 
 import mne
 import pandas as pd
 import yaml
 
 from moabb import paradigms as moabb_paradigms
 from moabb.analysis import analyze
+from moabb.datasets.base import BaseDataset
 from moabb.evaluations import (
     CrossSessionEvaluation,
     CrossSubjectEvaluation,
@@ -19,6 +21,7 @@
     generate_param_grid,
     parse_pipelines_from_directory,
 )
+from moabb.utils import set_download_dir
 
 
 try:
@@ -32,18 +35,19 @@
 
 
 def benchmark(  # noqa: C901
-    pipelines="./pipelines/",
-    evaluations=None,
-    paradigms=None,
-    results="./results/",
-    overwrite=False,
-    output="./benchmark/",
-    n_jobs=-1,
-    n_jobs_evaluation=1,
-    plot=False,
-    contexts=None,
-    include_datasets=None,
-    exclude_datasets=None,
+    pipelines: str = "./pipelines/",
+    evaluations: List[str] = None,
+    paradigms: List[str] = None,
+    results: str = "./results/",
+    overwrite: bool = False,
+    output: str = "./benchmark/",
+    n_jobs: int = -1,
+    n_jobs_evaluation: int = 1,
+    plot: bool = False,
+    contexts: str = None,
+    include_datasets: List[Union[str, BaseDataset]] = None,
+    exclude_datasets: List[Union[str, BaseDataset]] = None,
+    mne_data: str = None,
 ):
     """Run benchmarks for selected pipelines and datasets
 
@@ -100,6 +104,8 @@ def benchmark(  # noqa: C901
         and exclude_datasets are specified, raise an error.
     exclude_datasets: list of str or Dataset object
         Datasets to exclude from the benchmark run
+    mne_data: str
+        Folder where to save and load the datasets with mne structure.
 
     Returns
     -------
@@ -123,6 +129,9 @@ def benchmark(  # noqa: C901
     mne.set_log_level(False)
     # logging.basicConfig(level=logging.WARNING)
 
+    if mne_data is not None:
+        set_download_dir(mne_data)
+
     output = Path(output)
     if not osp.isdir(output):
         os.makedirs(output)
@@ -324,3 +333,9 @@ def _inc_exc_datasets(datasets, include_datasets, exclude_datasets):
     else:
         d = list(datasets)
     return d
+
+
+if __name__ == "__main__":
+    import fire
+
+    fire.Fire(benchmark)

pyproject.toml

@@ -62,6 +62,11 @@ tdlda = { git = "https://github.com/jsosulski/tdlda.git", rev = "0.1.0" }
 sphinx-design = "^0.3.0"
 sphinx-rtd-theme = "^1.2.0"
 
+[tool.poetry.group.benchmark]
+optional = true
+
+[tool.poetry.group.benchmark.dependencies]
+fire = "^0.5.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]