Add remaining code for optimization experiments

Author: amorehead

.gitignore

@@ -168,7 +168,10 @@ scripts/nautilus/persistent_storage.yaml
 
 bio-diffusion/
 logs/
+*optim_mols/
 outputs/
+epoch_*/
+.cache/
 
 data/EDM/GEOM
 data/EDM/QM9

configs/experiment/geom_mol_gen_ddpm.yaml

@@ -22,10 +22,10 @@ callbacks:
     save_top_k: 3
   early_stopping:
     monitor: "val/loss"
-    patience: 10
+    patience: 20
 
 trainer:
-  min_epochs: 30
+  min_epochs: 50
   max_epochs: 3000
   strategy: ddp_find_unused_parameters_false
   accelerator: gpu
@@ -122,6 +122,6 @@ datamodule:
 
 logger:
   wandb:
-    name: 02092023_GEOMMoleculeGenerationDDPM
+    name: 04302023_GEOMMoleculeGenerationDDPM
     group: "GEOM"
     tags: ${tags}

configs/experiment/geom_mol_gen_ddpm_grid_search.yaml

@@ -22,10 +22,10 @@ callbacks:
     save_top_k: 3
   early_stopping:
     monitor: "val/loss"
-    patience: 10
+    patience: 20
 
 trainer:
-  min_epochs: 30
+  min_epochs: 50
   max_epochs: 3000
   strategy: ddp_find_unused_parameters_false
   accelerator: gpu
@@ -125,7 +125,7 @@ datamodule:
 
 logger:
   wandb:
-    name: 02092023_GEOMMoleculeGenerationDDPM
+    name: 04302023_GEOMMoleculeGenerationDDPM
     group: "GEOM"
     tags: ${tags}
 

configs/model/geom_mol_gen_ddpm.yaml

@@ -19,3 +19,5 @@ defaults:
   - module_cfg: geom_mol_gen_ddpm_gcp_module.yaml
   - layer_cfg: geom_mol_gen_ddpm_gcp_interaction_layer.yaml
   - diffusion_cfg: geom_mol_gen_ddpm.yaml
+
+seed: ${..seed}

configs/mol_gen_eval.yaml

@@ -19,6 +19,7 @@ tags: ["dev"]
 ckpt_path: ???
 
 # inference (i.e., sampling) and evaluation arguments
+seed: 42
 num_samples: 10000
 sampling_batch_size: 100
 num_test_passes: 5

configs/mol_gen_eval_conditional_qm9.yaml

@@ -16,6 +16,7 @@ task_name: "mol_gen_eval_conditional_qm9"
 tags: ["dev"]
 
 # inference (i.e., sampling) and evaluation arguments
+seed: 42
 generator_model_filepath: ???
 classifier_model_dir: ""
 property: "alpha"

configs/mol_gen_sample.yaml

@@ -18,6 +18,7 @@ tags: ["dev"]
 ckpt_path: ???
 
 # inference (i.e., sampling) arguments
+seed: ???
 output_dir: ""
 num_samples: ???
 num_nodes: 19

configs/trainer/default.yaml

@@ -2,7 +2,7 @@ _target_: pytorch_lightning.Trainer
 
 default_root_dir: ${paths.output_dir}
 
-min_epochs: 5  # prevents early stopping
+min_epochs: 50  # prevents early stopping
 max_epochs: 3000
 
 strategy: ddp_find_unused_parameters_false

scripts/nautilus/README.md

@@ -18,15 +18,16 @@
 
 from omegaconf import DictConfig
 
-from torchtyping import TensorType, patch_typeguard
-from typeguard import typechecked
 from src.datamodules.components.edm_dataset import _edge_features, _node_features
 
 from src.models.components import ScalarVector, centralize
 from src.models.components.gcpnet import GCPNetDynamics
 from src.models.components.variational_diffusion import NODE_FEATURE_DIFFUSION_TARGETS
 from src.utils.pylogger import get_pylogger
 
+from torchtyping import TensorType, patch_typeguard
+from typeguard import typechecked
+
 patch_typeguard()  # use before @typechecked
 
 

src/models/components/egnn.py

@@ -16,15 +16,16 @@
 
 from omegaconf import OmegaConf, DictConfig
 
-from torchtyping import TensorType, patch_typeguard
-from typeguard import typechecked
 from src.datamodules.components.edm_dataset import _edge_features, _node_features
 
 from src.models import get_nonlinearity
 from src.models.components import GCPDropout, GCPLayerNorm, ScalarVector, centralize, is_identity, localize, safe_norm, scalarize, vectorize
 from src.models.components.variational_diffusion import NODE_FEATURE_DIFFUSION_TARGETS
 from src.utils.pylogger import get_pylogger
 
+from torchtyping import TensorType, patch_typeguard
+from typeguard import typechecked
+
 patch_typeguard()  # use before @typechecked