Add kwarg verbose=False to train_ensemble() (#40)

* move results/ dir creation to save_results_dict()

previously in example notebooks
also delete runs/ dir creation, handled by TensorBoard automatically

* add kwarg verbose=False to train_ensemble()

Author: janosh

aviary/utils.py

@@ -293,6 +293,7 @@ def train_ensemble(
     model_params: dict[str, Any],
     loss_dict: dict[str, Literal["L1", "L2", "CSE"]],
     patience: int = None,
+    verbose: bool = False,
 ) -> None:
     """Convenience method to train multiple models in serial.
 
@@ -313,6 +314,7 @@ def train_ensemble(
             to loss functions.
         patience (int, optional): Maximum number of epochs without improvement
             when early stopping. Defaults to None.
+        verbose (bool, optional): Whether to show progress bars for each epoch.
     """
     train_generator = DataLoader(train_set, **data_params)
     print(f"Training on {len(train_set):,} samples")
@@ -359,9 +361,7 @@ def train_ensemble(
 
         if log:
             writer = SummaryWriter(
-                log_dir=(
-                    f"runs/{model_name}/{model_name}-r{j}_{datetime.now():%d-%m-%Y_%H-%M-%S}"
-                )
+                f"runs/{model_name}/{model_name}-r{j}_{datetime.now():%d-%m-%Y_%H-%M-%S}"
             )
         else:
             writer = None
@@ -375,7 +375,7 @@ def train_ensemble(
                     optimizer=None,
                     normalizer_dict=normalizer_dict,
                     action="val",
-                    verbose=True,
+                    verbose=verbose,
                 )
 
                 val_score = {}
@@ -727,7 +727,7 @@ def save_results_dict(
         for col, data in results_dict[target_name].items():
 
             # NOTE we save pre_logits rather than logits due to fact
-            # that with the hetroskedastic setup we want to be able to
+            # that with the heteroskedastic setup we want to be able to
             # sample from the Gaussian distributed pre_logits we parameterise.
             if "pre-logits" in col:
                 for n_ens, y_pre_logit in enumerate(data):
@@ -760,6 +760,8 @@ def save_results_dict(
 
     file_name = model_name.replace("/", "_")
 
+    os.makedirs("results", exist_ok=True)
+
     csv_path = f"results/{file_name}.csv"
     df.to_csv(csv_path, index=False)
     print(f"\nSaved model predictions to '{csv_path}'")

examples/cgcnn-example.py

@@ -202,11 +202,6 @@ def main(  # noqa: C901
         "n_hidden": n_hidden,
     }
 
-    if log:
-        os.makedirs("runs/", exist_ok=True)
-
-    os.makedirs("results/", exist_ok=True)
-
     if train:
         train_ensemble(
             model_class=CrystalGraphConvNet,

examples/colab/Roost.ipynb

@@ -283,8 +283,6 @@
       },
       "outputs": [],
       "source": [
-        "import os \n",
-        "import numpy as np\n",
         "import torch\n",
         "from sklearn.model_selection import train_test_split as split\n",
         "\n",
@@ -3874,9 +3872,6 @@
         "    \"out_hidden\": [64, 64],\n",
         "}\n",
         "\n",
-        "os.makedirs(f\"models/{model_name}\", exist_ok=True)\n",
-        "os.makedirs(f\"results/{model_name}\", exist_ok=True)\n",
-        "\n",
         "train_ensemble(\n",
         "    model_class=Roost,\n",
         "    model_name=model_name,\n",

examples/colab/Wren.ipynb

@@ -283,8 +283,6 @@
       },
       "outputs": [],
       "source": [
-        "import os \n",
-        "import numpy as np\n",
         "import torch\n",
         "from sklearn.model_selection import train_test_split as split\n",
         "\n",
@@ -3879,9 +3877,6 @@
         "    \"out_hidden\": [64, 64],\n",
         "}\n",
         "\n",
-        "os.makedirs(f\"models/{model_name}\", exist_ok=True)\n",
-        "os.makedirs(f\"results/{model_name}\", exist_ok=True)\n",
-        "\n",
         "train_ensemble(\n",
         "    model_class=Wren,\n",
         "    model_name=model_name,\n",

examples/roost-example.py

@@ -186,11 +186,6 @@ def main(  # noqa: C901
         "out_hidden": [256, 128, 64],
     }
 
-    if log:
-        os.makedirs("runs/", exist_ok=True)
-
-    os.makedirs("results/", exist_ok=True)
-
     # TODO dump all args/kwargs to a file for reproducibility.
 
     if train:

examples/wren-example.py

@@ -199,11 +199,6 @@ def main(  # noqa: C901
         "trunk_hidden": [128, 64],
     }
 
-    if log:
-        os.makedirs("runs/", exist_ok=True)
-
-    os.makedirs("results/", exist_ok=True)
-
     if train:
         train_ensemble(
             model_class=Wren,