Undo accidental training script change

train_openfold.py

@@ -8,6 +8,7 @@ import os
 #os.environ["NODE_RANK"]="0"
 
 import random
+import sys
 import time
 
 import numpy as np
@@ -26,14 +27,22 @@ from openfold.data.data_modules import (
 )
 from openfold.model.model import AlphaFold
 from openfold.model.torchscript import script_preset_
+from openfold.np import residue_constants
+from openfold.utils.argparse import remove_arguments
 from openfold.utils.callbacks import (
     EarlyStoppingVerbose,
 )
 from openfold.utils.exponential_moving_average import ExponentialMovingAverage
-from openfold.utils.argparse import remove_arguments
 from openfold.utils.loss import AlphaFoldLoss, lddt_ca
+from openfold.utils.lr_schedulers import AlphaFoldLRScheduler
 from openfold.utils.seed import seed_everything
+from openfold.utils.superimposition import superimpose
 from openfold.utils.tensor_utils import tensor_tree_map
+from openfold.utils.validation_metrics import (
+    drmsd,
+    gdt_ts,
+    gdt_ha,
+)
 from scripts.zero_to_fp32 import (
     get_fp32_state_dict_from_zero_checkpoint
 )
@@ -57,68 +66,173 @@ class OpenFoldWrapper(pl.LightningModule):
     def forward(self, batch):
         return self.model(batch)
 
+    def _log(self, loss_breakdown, batch, outputs, train=True):
+        phase = "train" if train else "val"
+        for loss_name, indiv_loss in loss_breakdown.items():
+            self.log(
+                f"{phase}/{loss_name}", 
+                indiv_loss, 
+                on_step=train, on_epoch=(not train), logger=True,
+            )
+
+            if(train):
+                self.log(
+                    f"{phase}/{loss_name}_epoch",
+                    indiv_loss,
+                    on_step=False, on_epoch=True, logger=True,
+                )
+
+        with torch.no_grad():
+            other_metrics = self._compute_validation_metrics(
+                batch, 
+                outputs,
+                superimposition_metrics=(not train)
+            )
+
+        for k,v in other_metrics.items():
+            self.log(
+                f"{phase}/{k}", 
+                v, 
+                on_step=False, on_epoch=True, logger=True
+            )
+
     def training_step(self, batch, batch_idx):
         if(self.ema.device != batch["aatype"].device):
             self.ema.to(batch["aatype"].device)
 
         # Run the model
         outputs = self(batch)
-        
+
         # Remove the recycling dimension
         batch = tensor_tree_map(lambda t: t[..., -1], batch)
 
         # Compute loss
-        loss = self.loss(outputs, batch)
+        loss, loss_breakdown = self.loss(
+            outputs, batch, _return_breakdown=True
+        )
 
-        self.log("train/loss", loss, on_step=True, logger=True)
+        # Log it
+        self._log(loss_breakdown, batch, outputs)
 
         return loss
 
-    def training_step_end(self, outputs):
-        # Temporary measure to address DeepSpeed scheduler bug
-        if(self.trainer.global_step != self.last_lr_step):
-            self.lr_schedulers().step()
-            self.last_lr_step = self.trainer.global_step
+    def on_before_zero_grad(self, *args, **kwargs):
+        self.ema.update(self.model)
 
     def validation_step(self, batch, batch_idx):
         # At the start of validation, load the EMA weights
         if(self.cached_weights is None):
-            self.cached_weights = self.model.state_dict()
+            # model.state_dict() contains references to model weights rather
+            # than copies. Therefore, we need to clone them before calling 
+            # load_state_dict().
+            clone_param = lambda t: t.detach().clone()
+            self.cached_weights = tensor_tree_map(clone_param, self.model.state_dict())
             self.model.load_state_dict(self.ema.state_dict()["params"])
        
-        # Calculate validation loss
+        # Run the model
         outputs = self(batch)
         batch = tensor_tree_map(lambda t: t[..., -1], batch)
-        loss = lddt_ca(
-            outputs["final_atom_positions"],
-            batch["all_atom_positions"],
-            batch["all_atom_mask"],
-            eps=self.config.globals.eps,
-            per_residue=False,
+
+        # Compute loss and other metrics
+        batch["use_clamped_fape"] = 0.
+        _, loss_breakdown = self.loss(
+            outputs, batch, _return_breakdown=True
         )
-        self.log("val/loss", loss, logger=True)
 
+        self._log(loss_breakdown, batch, outputs, train=False)
+        
     def validation_epoch_end(self, _):
         # Restore the model weights to normal
         self.model.load_state_dict(self.cached_weights)
         self.cached_weights = None
 
+    def _compute_validation_metrics(self, 
+        batch, 
+        outputs, 
+        superimposition_metrics=False
+    ):
+        metrics = {}
+        
+        gt_coords = batch["all_atom_positions"]
+        pred_coords = outputs["final_atom_positions"]
+        all_atom_mask = batch["all_atom_mask"]
+    
+        # This is super janky for superimposition. Fix later
+        gt_coords_masked = gt_coords * all_atom_mask[..., None]
+        pred_coords_masked = pred_coords * all_atom_mask[..., None]
+        ca_pos = residue_constants.atom_order["CA"]
+        gt_coords_masked_ca = gt_coords_masked[..., ca_pos, :]
+        pred_coords_masked_ca = pred_coords_masked[..., ca_pos, :]
+        all_atom_mask_ca = all_atom_mask[..., ca_pos]
+    
+        lddt_ca_score = lddt_ca(
+            pred_coords,
+            gt_coords,
+            all_atom_mask,
+            eps=self.config.globals.eps,
+            per_residue=False,
+        )
+   
+        metrics["lddt_ca"] = lddt_ca_score
+   
+        drmsd_ca_score = drmsd(
+            pred_coords_masked_ca,
+            gt_coords_masked_ca,
+            mask=all_atom_mask_ca, # still required here to compute n
+        )
+   
+        metrics["drmsd_ca"] = drmsd_ca_score
+    
+        if(superimposition_metrics):
+            superimposed_pred, alignment_rmsd = superimpose(
+                gt_coords_masked_ca, pred_coords_masked_ca, all_atom_mask_ca,
+            )
+            gdt_ts_score = gdt_ts(
+                superimposed_pred, gt_coords_masked_ca, all_atom_mask_ca
+            )
+            gdt_ha_score = gdt_ha(
+                superimposed_pred, gt_coords_masked_ca, all_atom_mask_ca
+            )
+
+            metrics["alignment_rmsd"] = alignment_rmsd
+            metrics["gdt_ts"] = gdt_ts_score
+            metrics["gdt_ha"] = gdt_ha_score
+    
+        return metrics
+
     def configure_optimizers(self, 
         learning_rate: float = 1e-3,
         eps: float = 1e-5,
     ) -> torch.optim.Adam:
+#        return torch.optim.Adam(
+#            self.model.parameters(),
+#            lr=learning_rate,
+#            eps=eps
+#        )
         # Ignored as long as a DeepSpeed optimizer is configured
-        return torch.optim.Adam(
+        optimizer = torch.optim.Adam(
             self.model.parameters(), 
             lr=learning_rate, 
             eps=eps
         )
+        lr_scheduler = AlphaFoldLRScheduler(
+            optimizer,
+        )
 
-    def on_before_zero_grad(self, *args, **kwargs):
-        self.ema.update(self.model)
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": {
+                "scheduler": lr_scheduler,
+                "interval": "step",
+                "name": "AlphaFoldLRScheduler",
+            }
+        }
 
     def on_load_checkpoint(self, checkpoint):
-        self.ema.load_state_dict(checkpoint["ema"])
+        ema = checkpoint["ema"]
+        if(not self.model.template_config.enabled):
+            ema["params"] = {k:v for k,v in ema["params"].items() if not "template" in k}
+        self.ema.load_state_dict(ema)
 
     def on_save_checkpoint(self, checkpoint):
         checkpoint["ema"] = self.ema.state_dict()
@@ -129,7 +243,7 @@ def main(args):
         seed_everything(args.seed) 
 
     config = model_config(
-        "model_1", 
+        args.config_preset, 
         train=True, 
         low_prec=(args.precision == "16")
     ) 
@@ -140,7 +254,7 @@ def main(args):
         sd = {k[len("module."):]:v for k,v in sd.items()}
         model_module.load_state_dict(sd)
         logging.info("Successfully loaded model weights...")
-
+ 
     # TorchScript components of the model
     if(args.script_modules):
         script_preset_(model_module)
@@ -159,16 +273,18 @@ def main(args):
     if(args.checkpoint_every_epoch):
         mc = ModelCheckpoint(
             every_n_epochs=1,
+            auto_insert_metric_name=False,
+            save_top_k=-1,
         )
         callbacks.append(mc)
 
     if(args.early_stopping):
         es = EarlyStoppingVerbose(
-            monitor="val/loss",
+            monitor="val/lddt_ca",
             min_delta=args.min_delta,
             patience=args.patience,
             verbose=False,
-            mode="min",
+            mode="max",
             check_finite=True,
             strict=True,
         )
@@ -198,14 +314,8 @@ def main(args):
         loggers.append(wdb_logger)
 
     if(args.deepspeed_config_path is not None):
-        #if "SLURM_JOB_ID" in os.environ:
-        #    cluster_environment = SLURMEnvironment()
-        #else:
-        #    cluster_environment = None
-      
         strategy = DeepSpeedPlugin(
             config=args.deepspeed_config_path,
-        #    cluster_environment=cluster_environment,
         )
         if(args.wandb):
             wdb_logger.experiment.save(args.deepspeed_config_path)
@@ -214,7 +324,12 @@ def main(args):
         strategy = DDPPlugin(find_unused_parameters=False)
     else:
         strategy = None
-    
+ 
+    if(args.wandb):
+        freeze_path = f"{wdb_logger.experiment.dir}/package_versions.txt"
+        os.system(f"{sys.executable} -m pip freeze > {freeze_path}")
+        wdb_logger.experiment.save(f"{freeze_path}")
+
     trainer = pl.Trainer.from_argparse_args(
         args,
         default_root_dir=args.output_dir,
@@ -352,37 +467,65 @@ if __name__ == "__main__":
     )
     parser.add_argument(
         "--wandb", action="store_true", default=False,
+        help="Whether to log metrics to Weights & Biases"
     )
     parser.add_argument(
         "--experiment_name", type=str, default=None,
+        help="Name of the current experiment. Used for wandb logging"
     )
     parser.add_argument(
         "--wandb_id", type=str, default=None,
+        help="ID of a previous run to be resumed"
     )
     parser.add_argument(
         "--wandb_project", type=str, default=None,
+        help="Name of the wandb project to which this run will belong"
     )
     parser.add_argument(
         "--wandb_entity", type=str, default=None,
+        help="wandb username or team name to which runs are attributed"
     )
     parser.add_argument(
         "--script_modules", type=bool_type, default=False,
         help="Whether to TorchScript eligible components of them model"
     )
     parser.add_argument(
-        "--train_prot_data_cache_path", type=str, default=None,
+        "--train_chain_data_cache_path", type=str, default=None,
     )
     parser.add_argument(
-        "--distillation_prot_data_cache_path", type=str, default=None,
+        "--distillation_chain_data_cache_path", type=str, default=None,
     )
     parser.add_argument(
         "--train_epoch_len", type=int, default=10000,
+        help=(
+            "The virtual length of each training epoch. Stochastic filtering "
+            "of training data means that training datasets have no "
+            "well-defined length. This virtual length affects frequency of "
+            "validation & checkpointing (by default, one of each per epoch)."
+        )
     )
     parser.add_argument(
-        "--_alignment_index_path", type=str, default=None,
+        "--log_lr", action="store_true", default=False,
+        help="Whether to log the actual learning rate"
     )
     parser.add_argument(
-        "--log_lr", action="store_true", default=False,
+        "--config_preset", type=str, default="initial_training",
+        help=(
+            'Config setting. Choose e.g. "initial_training", "finetuning", '
+            '"model_1", etc. By default, the actual values in the config are '
+            'used.'
+        )
+    )
+    parser.add_argument(
+        "--_distillation_structure_index_path", type=str, default=None,
+    )
+    parser.add_argument(
+        "--alignment_index_path", type=str, default=None,
+        help="Training alignment index. See the README for instructions."
+    )
+    parser.add_argument(
+        "--distillation_alignment_index_path", type=str, default=None,
+        help="Distillation alignment index. See the README for instructions."
     )
     parser = pl.Trainer.add_argparse_args(parser)