Add wandb-sweep example and clean-up use of configs (#118)

* feat: added support for wandb logging (single runs, not sweeps yet)

* feat: replaced hps provided as a dict by a Config() object in seh_frag_moo

* feat: added config.desc

* fix: added train/valid for wandb log

* fix: allow JSON serialization of Enum objects

* chore: tox

* chore: replaced hps (dict) by Config() in all tasks. Moved qm9.py out of qm9/

* fix: changed default focus_region for frag_moo

* fix: added assert to prevent inadvertently manipulating a Config rather than Config() object

* removed cfg.use_wandb and simply test if wandb has been initialised in trainer

* chore: adding comment

* fix: typo

* chore: added cfg.task.seh_moo.log_topk to de-clutter a bit

* fix: added wandb to dependencies

* minor: file name change for consistency

* chore: centralised self.cfg.overwrite_existing_exp in GFNTrainer() (removed from all tasks to simplify mains)

* feat: added hyperopt/wandb_demo

* feat: removed wandb_agent_main.py to have the search and entrypoint defined in a single file

* chore: tox

* fix: minor in wandb_demo

* fix: storage path

* chore: tox

src/gflownet/algo/config.py

@@ -3,7 +3,7 @@
 from typing import Optional
 
 
-class TBVariant(Enum):
+class TBVariant(int, Enum):
     """See algo.trajectory_balance.TrajectoryBalance for details."""
 
     TB = 0
@@ -116,8 +116,6 @@ class AlgoConfig:
         Do not take random actions after this number of steps
     valid_random_action_prob : float
         The probability of taking a random action during validation
-    valid_sample_cond_info : bool
-        Whether to sample conditioning information during validation (if False, expects a validation set of cond_info)
     sampling_tau : float
         The EMA factor for the sampling model (theta_sampler = tau * theta_sampler + (1-tau) * theta)
     """
@@ -133,7 +131,6 @@ class AlgoConfig:
     train_random_action_prob: float = 0.0
     train_det_after: Optional[int] = None
     valid_random_action_prob: float = 0.0
-    valid_sample_cond_info: bool = True
     sampling_tau: float = 0.0
     tb: TBConfig = TBConfig()
     moql: MOQLConfig = MOQLConfig()

src/gflownet/config.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, fields, is_dataclass
 from typing import Optional
 
 from omegaconf import MISSING
@@ -50,6 +50,8 @@ class Config:
 
     Attributes
     ----------
+    desc : str
+        A description of the experiment
     log_dir : str
         The directory where to store logs, checkpoints, and samples.
     device : str
@@ -82,6 +84,7 @@ class Config:
         Whether to overwrite the contents of the log_dir if it already exists
     """
 
+    desc: str = "noDesc"
     log_dir: str = MISSING
     device: str = "cuda"
     seed: int = 0
@@ -96,10 +99,28 @@ class Config:
     hostname: Optional[str] = None
     pickle_mp_messages: bool = False
     git_hash: Optional[str] = None
-    overwrite_existing_exp: bool = True
+    overwrite_existing_exp: bool = False
     algo: AlgoConfig = AlgoConfig()
     model: ModelConfig = ModelConfig()
     opt: OptimizerConfig = OptimizerConfig()
     replay: ReplayConfig = ReplayConfig()
     task: TasksConfig = TasksConfig()
     cond: ConditionalsConfig = ConditionalsConfig()
+
+
+def init_empty(cfg):
+    """
+    Initialize a dataclass instance with all fields set to MISSING,
+    including nested dataclasses.
+
+    This is meant to be used on the user side (tasks) to provide
+    some configuration using the Config class while overwritting
+    only the fields that have been set by the user.
+    """
+    for f in fields(cfg):
+        if is_dataclass(f.type):
+            setattr(cfg, f.name, init_empty(f.type()))
+        else:
+            setattr(cfg, f.name, MISSING)
+
+    return cfg

src/gflownet/hyperopt/wandb_demo/README.md

@@ -0,0 +1,6 @@
+Everything is contained in one file; `init_wandb_sweep.py` both defines the search space of the sweep and is the entrypoint of wandb agents.
+
+To launch the search:
+1. `python init_wandb_sweep.py` to intialize the sweep
+2. `sbatch launch_wandb_agents.sh <SWEEP_ID>` to schedule a jobarray in slurm which will launch wandb agents.
+The number of jobs in the sbatch file should reflect the size of the hyperparameter space that is being sweeped.

src/gflownet/hyperopt/wandb_demo/init_wandb_sweep.py

@@ -0,0 +1,73 @@
+import os
+import sys
+import time
+
+import wandb
+
+from gflownet.config import Config, init_empty
+from gflownet.tasks.seh_frag_moo import SEHMOOFragTrainer
+
+TIME = time.strftime("%m-%d-%H-%M")
+ENTITY = "valencelabs"
+PROJECT = "gflownet"
+SWEEP_NAME = f"{TIME}-sehFragMoo-Zlr-Zlrdecay"
+STORAGE_DIR = f"~/storage/wandb_sweeps/{SWEEP_NAME}"
+
+
+# Define the search space of the sweep
+sweep_config = {
+    "name": SWEEP_NAME,
+    "program": "init_wandb_sweep.py",
+    "controller": {
+        "type": "cloud",
+    },
+    "method": "grid",
+    "parameters": {
+        "config.algo.tb.Z_learning_rate": {"values": [1e-4, 1e-3, 1e-2]},
+        "config.algo.tb.Z_lr_decay": {"values": [2_000, 50_000]},
+    },
+}
+
+
+def wandb_config_merger():
+    config = init_empty(Config())
+    wandb_config = wandb.config
+
+    # Set desired config values
+    config.log_dir = f"{STORAGE_DIR}/{wandb.run.name}-id-{wandb.run.id}"
+    config.print_every = 100
+    config.validate_every = 1000
+    config.num_final_gen_steps = 1000
+    config.num_training_steps = 40_000
+    config.pickle_mp_messages = True
+    config.overwrite_existing_exp = False
+    config.algo.sampling_tau = 0.95
+    config.algo.train_random_action_prob = 0.01
+    config.algo.tb.Z_learning_rate = 1e-3
+    config.task.seh_moo.objectives = ["seh", "qed"]
+    config.cond.temperature.sample_dist = "constant"
+    config.cond.temperature.dist_params = [60.0]
+    config.cond.weighted_prefs.preference_type = "dirichlet"
+    config.cond.focus_region.focus_type = None
+    config.replay.use = False
+
+    # Merge the wandb sweep config with the nested config from gflownet
+    config.algo.tb.Z_learning_rate = wandb_config["config.algo.tb.Z_learning_rate"]
+    config.algo.tb.Z_lr_decay = wandb_config["config.algo.tb.Z_lr_decay"]
+
+    return config
+
+
+if __name__ == "__main__":
+    # if there no arguments, initialize the sweep, otherwise this is a wandb agent
+    if len(sys.argv) == 1:
+        if os.path.exists(STORAGE_DIR):
+            raise ValueError(f"Sweep storage directory {STORAGE_DIR} already exists.")
+
+        wandb.sweep(sweep_config, entity=ENTITY, project=PROJECT)
+
+    else:
+        wandb.init(entity=ENTITY, project=PROJECT)
+        config = wandb_config_merger()
+        trial = SEHMOOFragTrainer(config)
+        trial.run()

src/gflownet/hyperopt/wandb_demo/launch_wandb_agents.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+# Purpose: Script to allocate a node and run a wandb sweep agent on it
+# Usage: sbatch launch_wandb_agent.sh <SWEEP_ID>
+
+#SBATCH --job-name=wandb_sweep_agent
+#SBATCH --array=1-6
+#SBATCH --time=23:59:00
+#SBATCH --output=slurm_output_files/%x_%N_%A_%a.out
+#SBATCH --gpus=1
+#SBATCH --cpus-per-task=16
+#SBATCH --mem=16GB
+#SBATCH --partition compute
+
+source activate gfn-py39-torch113
+echo "Using environment={$CONDA_DEFAULT_ENV}"
+
+# launch wandb agent
+wandb agent --count 1 --entity valencelabs --project gflownet $1

src/gflownet/models/config.py

@@ -9,7 +9,7 @@ class GraphTransformerConfig:
     num_mlp_layers: int = 0
 
 
-class SeqPosEnc(Enum):
+class SeqPosEnc(int, Enum):
     Pos = 0
     Rotary = 1
 

src/gflownet/online_trainer.py

@@ -105,13 +105,13 @@ def setup(self):
         git_hash = git.Repo(__file__, search_parent_directories=True).head.object.hexsha[:7]
         self.cfg.git_hash = git_hash
 
-        yaml = OmegaConf.to_yaml(self.cfg)
-        os.makedirs(self.cfg.log_dir, exist_ok=True)
-        if self.print_hps:
+        yaml_cfg = OmegaConf.to_yaml(self.cfg)
+        if self.print_config:
             print("\n\nHyperparameters:\n")
-            print(yaml)
-        with open(pathlib.Path(self.cfg.log_dir) / "hps.yaml", "w", encoding="utf8") as f:
-            f.write(yaml)
+            print(yaml_cfg)
+        os.makedirs(self.cfg.log_dir, exist_ok=True)
+        with open(pathlib.Path(self.cfg.log_dir) / "config.yaml", "w", encoding="utf8") as f:
+            f.write(yaml_cfg)
 
     def step(self, loss: Tensor):
         loss.backward()

src/gflownet/tasks/config.py

@@ -26,6 +26,7 @@ class SEHMOOTaskConfig:
     n_valid: int = 15
     n_valid_repeats: int = 128
     objectives: List[str] = field(default_factory=lambda: ["seh", "qed", "sa", "mw"])
+    log_topk: bool = False
     online_pareto_front: bool = True
 
 

src/gflownet/tasks/make_rings.py

@@ -1,4 +1,3 @@
-import os
 import socket
 from typing import Dict, List, Tuple, Union
 
@@ -8,7 +7,7 @@
 from rdkit.Chem.rdchem import Mol as RDMol
 from torch import Tensor
 
-from gflownet.config import Config
+from gflownet.config import Config, init_empty
 from gflownet.envs.mol_building_env import MolBuildingEnvContext
 from gflownet.online_trainer import StandardOnlineTrainer
 from gflownet.trainer import FlatRewards, GFNTask, RewardScalar
@@ -72,17 +71,16 @@ def setup_env_context(self):
 
 
 def main():
-    hps = {
-        "log_dir": "./logs/debug_run_mr4",
-        "device": "cuda",
-        "num_training_steps": 10_000,
-        "num_workers": 8,
-        "algo": {"tb": {"do_parameterize_p_b": True}},
-    }
-    os.makedirs(hps["log_dir"], exist_ok=True)
-
-    trial = MakeRingsTrainer(hps)
-    trial.print_every = 1
+    """Example of how this model can be run."""
+    config = init_empty(Config())
+    config.print_every = 1
+    config.log_dir = "./logs/debug_run_mr4"
+    config.device = "cuda"
+    config.num_training_steps = 10_000
+    config.num_workers = 8
+    config.algo.tb.do_parameterize_p_b = True
+
+    trial = MakeRingsTrainer(config)
     trial.run()
 
 

src/gflownet/tasks/qm9/qm9.py → src/gflownet/tasks/qm9.py

@@ -9,7 +9,7 @@
 from torch.utils.data import Dataset
 
 import gflownet.models.mxmnet as mxmnet
-from gflownet.config import Config
+from gflownet.config import Config, init_empty
 from gflownet.data.qm9 import QM9Dataset
 from gflownet.envs.mol_building_env import MolBuildingEnvContext
 from gflownet.online_trainer import StandardOnlineTrainer
@@ -162,3 +162,22 @@ def setup(self):
         super().setup()
         self.training_data.setup(self.task, self.ctx)
         self.test_data.setup(self.task, self.ctx)
+
+
+def main():
+    """Example of how this model can be run."""
+    config = init_empty(Config())
+    config.num_workers = 0
+    config.num_training_steps = 100000
+    config.validate_every = 100
+    config.log_dir = "./logs/debug_qm9"
+    config.opt.lr_decay = 10000
+    config.task.qm9.h5_path = "/rxrx/data/chem/qm9/qm9.h5"
+    config.task.qm9.model_path = "/rxrx/data/chem/qm9/mxmnet_gap_model.pt"
+
+    trial = QM9GapTrainer(config)
+    trial.run()
+
+
+if __name__ == "__main__":
+    main()

src/gflownet/tasks/qm9/qm9_moo.py → src/gflownet/tasks/qm9_moo.py

@@ -15,7 +15,7 @@
 from gflownet.config import Config
 from gflownet.data.qm9 import QM9Dataset
 from gflownet.envs.mol_building_env import MolBuildingEnvContext
-from gflownet.tasks.qm9.qm9 import QM9GapTask, QM9GapTrainer
+from gflownet.tasks.qm9 import QM9GapTask, QM9GapTrainer
 from gflownet.tasks.seh_frag_moo import RepeatedCondInfoDataset, aux_tasks
 from gflownet.trainer import FlatRewards, RewardScalar
 from gflownet.utils import metrics
@@ -197,7 +197,7 @@ def set_default_hps(self, cfg: Config):
         cfg.algo.sampling_tau = 0.95
         # We use a fixed set of preferences as our "validation set", so we must disable the preference (cond_info)
         # sampling and set the offline ratio to 1
-        cfg.algo.valid_sample_cond_info = False
+        cfg.cond.valid_sample_cond_info = False
         cfg.algo.valid_offline_ratio = 1
 
     def setup_algo(self):

src/gflownet/tasks/seh_frag.py

@@ -1,5 +1,3 @@
-import os
-import shutil
 import socket
 from typing import Callable, Dict, List, Tuple, Union
 
@@ -13,7 +11,7 @@
 from torch.utils.data import Dataset
 from torch_geometric.data import Data
 
-from gflownet.config import Config
+from gflownet.config import Config, init_empty
 from gflownet.envs.frag_mol_env import FragMolBuildingEnvContext, Graph
 from gflownet.models import bengio2021flow
 from gflownet.online_trainer import StandardOnlineTrainer
@@ -200,33 +198,21 @@ def setup(self):
 
 
 def main():
-    """Example of how this trainer can be run"""
-    hps = {
-        "log_dir": "./logs/debug_run_seh_frag_pb",
-        "device": "cuda" if torch.cuda.is_available() else "cpu",
-        "overwrite_existing_exp": True,
-        "num_training_steps": 10_000,
-        "num_workers": 8,
-        "opt": {
-            "lr_decay": 20000,
-        },
-        "algo": {"sampling_tau": 0.99, "offline_ratio": 0.0},
-        "cond": {
-            "temperature": {
-                "sample_dist": "uniform",
-                "dist_params": [0, 64.0],
-            }
-        },
-    }
-    if os.path.exists(hps["log_dir"]):
-        if hps["overwrite_existing_exp"]:
-            shutil.rmtree(hps["log_dir"])
-        else:
-            raise ValueError(f"Log dir {hps['log_dir']} already exists. Set overwrite_existing_exp=True to delete it.")
-    os.makedirs(hps["log_dir"])
-
-    trial = SEHFragTrainer(hps)
-    trial.print_every = 1
+    """Example of how this model can be run."""
+    config = init_empty(Config())
+    config.print_every = 1
+    config.log_dir = "./logs/debug_run_seh_frag_pb"
+    config.device = "cuda" if torch.cuda.is_available() else "cpu"
+    config.overwrite_existing_exp = True
+    config.num_training_steps = 10_000
+    config.num_workers = 8
+    config.opt.lr_decay = 20_000
+    config.algo.sampling_tau = 0.99
+    config.algo.offline_ratio = 0.0
+    config.cond.temperature.sample_dist = "uniform"
+    config.cond.temperature.dist_params = [0, 64.0]
+
+    trial = SEHFragTrainer(config)
     trial.run()