Faster tests (#579)

Author: Luthaf

.github/workflows/architecture-tests.yml

@@ -8,6 +8,7 @@ on:
 
 jobs:
   tests:
+    name: ${{ matrix.architecture-name }}
     strategy:
       matrix:
         include:

tests/cli/__init__.py

@@ -15,3 +15,4 @@
 MODEL_PATH = RESOURCES_PATH / "model-32-bit.pt"
 MODEL_PATH_64_BIT = RESOURCES_PATH / "model-64-bit.ckpt"
 OPTIONS_PATH = RESOURCES_PATH / "options.yaml"
+OPTIONS_NANOPET_PATH = RESOURCES_PATH / "options-nanopet.yaml"

tests/cli/test_train_model.py

@@ -27,6 +27,7 @@
     DATASET_PATH_QM7X,
     DATASET_PATH_QM9,
     MODEL_PATH_64_BIT,
+    OPTIONS_NANOPET_PATH,
     OPTIONS_PATH,
     RESOURCES_PATH,
 )
@@ -114,7 +115,7 @@ def test_train(capfd, monkeypatch, tmp_path, output):
     assert "train" in stdout_log
     assert "energy" in stdout_log
     assert "with index" not in stdout_log  # index only printed for more than 1 dataset
-    assert "Running final evaluation with batch size 2" in stdout_log
+    assert "Running final evaluation with batch size 16" in stdout_log
 
     # Open the CSV log file and check if the logging is correct
     csv_glob = glob.glob("outputs/*/*/train.csv")
@@ -667,8 +668,7 @@ def test_train_generic_target_metatensor(monkeypatch, tmp_path, with_scalar_part
     )
 
     # run training with original options
-    options = OmegaConf.load(OPTIONS_PATH)
-    options["architecture"]["name"] = "experimental.nanopet"
+    options = OmegaConf.load(OPTIONS_NANOPET_PATH)
     options["training_set"]["systems"]["read_from"] = "qm7x_reduced_100.xyz"
     options["training_set"]["targets"] = {
         "mtt::polarizability": {

tests/resources/generate-outputs.sh

@@ -9,7 +9,7 @@ cd $ROOT_DIR
 
 mtt train options.yaml -o model-32-bit.pt -r base_precision=32 # > /dev/null
 mtt train options.yaml -o model-64-bit.pt -r base_precision=64 # > /dev/null
-mtt train options.yaml -o model-no-extensions.pt -r architecture.name=experimental.nanopet # > /dev/null
+mtt train options-nanopet.yaml -o model-no-extensions.pt # > /dev/null
 
 # upload results to private HF repo if token is set
 if [ -n "${HUGGINGFACE_TOKEN_METATRAIN:-}" ]; then

tests/resources/options-nanopet.yaml

@@ -0,0 +1,19 @@
+seed: 42
+
+architecture:
+  name: experimental.nanopet
+  training:
+    batch_size: 16
+    num_epochs: 1
+
+training_set:
+  systems:
+    read_from: qm9_reduced_100.xyz
+    length_unit: angstrom
+  targets:
+    energy:
+      key: U0
+      unit: eV
+
+test_set: 0.5
+validation_set: 0.1

tests/resources/options.yaml

@@ -3,8 +3,12 @@ seed: 42
 architecture:
   name: soap_bpnn
   training:
-    batch_size: 2
+    batch_size: 16
     num_epochs: 1
+  model:
+    soap:
+      max_radial: 4
+      max_angular: 2
 
 training_set:
   systems: