fix bad import in train_cgcnn

from {aviary->matbench_discovery} import ROOT

Author: janosh

models/bowsr/test_bowsr.py

@@ -106,7 +106,6 @@
 
 slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
 wandb.init(
-    entity="janosh",
     project="matbench-discovery",
     name=f"{job_name}-{slurm_job_id}-{slurm_array_task_id}",
     config=run_params,

models/cgcnn/test_cgcnn.py

@@ -27,8 +27,8 @@
 stores predictions to CSV.
 """
 
-module_dir = os.path.dirname(__file__)
 today = f"{datetime.now():%Y-%m-%d}"
+log_dir = f"{os.path.dirname(__file__)}/{today}-test"
 ensemble_id = "cgcnn-e_form-ensemble-1"
 run_name = f"{ensemble_id}-IS2RE"
 
@@ -37,7 +37,7 @@
     partition="ampere",
     account="LEE-SL3-GPU",
     time="1:0:0",
-    log_dir=module_dir,
+    log_dir=log_dir,
     slurm_flags=("--nodes", "1", "--gpus-per-node", "1"),
 )
 
@@ -82,7 +82,7 @@
     data_loader=data_loader,
 )
 
-df.round(6).to_csv(f"{module_dir}/{today}-{run_name}-preds.csv", index=False)
+df.round(6).to_csv(f"{log_dir}/{today}-{run_name}-preds.csv", index=False)
 
 
 # %%

models/cgcnn/train_cgcnn.py

@@ -3,7 +3,6 @@
 from datetime import datetime
 
 import pandas as pd
-from aviary import ROOT
 from aviary.cgcnn.data import CrystalGraphData, collate_batch
 from aviary.cgcnn.model import CrystalGraphConvNet
 from aviary.core import TaskType
@@ -12,10 +11,11 @@
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
+from matbench_discovery import ROOT
 from matbench_discovery.slurm import slurm_submit
 
 """
-Train a CGCNN ensemble of size n_folds on target_col of data_path.
+Train a CGCNN ensemble of size n_ens on target_col of data_path.
 """
 
 __author__ = "Janosh Riebesell"
@@ -28,7 +28,7 @@
 run_name = f"cgcnn-robust-{target_col}"
 print(f"{run_name=}")
 robust = "robust" in run_name.lower()
-n_folds = 10
+n_ens = 10
 timestamp = f"{datetime.now():%Y-%m-%d@%H-%M-%S}"
 today = timestamp.split("@")[0]
 log_dir = f"{os.path.dirname(__file__)}/{today}-{run_name}"
@@ -38,7 +38,7 @@
     partition="ampere",
     account="LEE-SL3-GPU",
     time="8:0:0",
-    array=f"1-{n_folds}",
+    array=f"1-{n_ens}",
     log_dir=log_dir,
     slurm_flags=("--nodes", "1", "--gpus-per-node", "1"),
 )

models/voronoi/voronoi_featurize_dataset.py

@@ -15,7 +15,7 @@
 
 today = f"{datetime.now():%Y-%m-%d}"
 module_dir = os.path.dirname(__file__)
-
+assert featurizer._n_jobs == 1, "set n_jobs=1 to avoid OOM errors"
 
 data_name = "mp"  # "mp"
 if data_name == "wbm":
@@ -35,15 +35,14 @@
     account="LEE-SL3-CPU",
     time=(slurm_max_job_time := "8:0:0"),
     array=f"1-{slurm_array_task_count}",
+    slurm_flags=("--mem", "30G") if data_name == "mp" else (),
     log_dir=log_dir,
 )
 
 
 # %%
 slurm_array_task_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))
-slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
-run_name = f"{job_name}-{slurm_job_id}-{slurm_array_task_id}"
-out_path = f"{log_dir}/{run_name}.csv.bz2"
+out_path = f"{log_dir}/{job_name}.csv.bz2"
 
 if os.path.isfile(out_path):
     raise SystemExit(f"{out_path = } already exists, exciting early")
@@ -75,9 +74,10 @@
 if wandb.run is None:
     wandb.login()
 
+slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
 wandb.init(
     project="matbench-discovery",
-    name=run_name,
+    name=f"{job_name}-{slurm_job_id}-{slurm_array_task_id}",
     config=run_params,
 )
 

models/wrenformer/train_wrenformer.py

@@ -9,7 +9,7 @@
 from matbench_discovery.slurm import slurm_submit
 
 """
-Train a Wrenformer ensemble of size n_folds on target_col of data_path.
+Train a Wrenformer ensemble of size n_ens on target_col of data_path.
 """
 
 __author__ = "Janosh Riebesell"
@@ -25,7 +25,7 @@
 # target_col = "mp_energy_per_atom"
 data_name = "m3gnet-trainset" if "m3gnet" in data_path else "mp"
 run_name = f"train-wrenformer-robust-{data_name}-{target_col}"
-n_folds = 10
+n_ens = 10
 timestamp = f"{datetime.now():%Y-%m-%d@%H-%M-%S}"
 today = timestamp.split("@")[0]
 dataset = "mp"
@@ -36,7 +36,7 @@
     partition="ampere",
     account="LEE-SL3-GPU",
     time="8:0:0",
-    array=f"1-{n_folds}",
+    array=f"1-{n_ens}",
     log_dir=log_dir,
     slurm_flags=("--nodes", "1", "--gpus-per-node", "1"),
 )
@@ -73,7 +73,7 @@
     target_col=target_col,
     task_type="regression",
     timestamp=timestamp,
-    # folds=(n_folds, slurm_array_task_id),
+    # folds=(n_ens, slurm_array_task_id),
     epochs=epochs,
     checkpoint="wandb",  # None | 'local' | 'wandb',
     input_col=input_col,