train.py

import json
import sys
from typing import Optional

# This import must be on top to set the environment variables before importing other modules
import env_consts
import time
import os

from lightning.pytorch import seed_everything
import lightning.pytorch as pl
import torch
from lightning.pytorch.callbacks import LearningRateMonitor, ModelCheckpoint
from lightning.pytorch.loggers import WandbLogger
from lightning.pytorch.profilers import AdvancedProfiler

from dockformer.config import model_config
from dockformer.data.data_modules import OpenFoldDataModule, DockFormerDataModule
from dockformer.model.model import AlphaFold
from dockformer.utils import residue_constants
from dockformer.utils.exponential_moving_average import ExponentialMovingAverage
from dockformer.utils.loss import AlphaFoldLoss, lddt_ca
from dockformer.utils.lr_schedulers import AlphaFoldLRScheduler
from dockformer.utils.script_utils import get_latest_checkpoint
from dockformer.utils.superimposition import superimpose
from dockformer.utils.tensor_utils import tensor_tree_map
from dockformer.utils.validation_metrics import (
    drmsd,
    gdt_ts,
    gdt_ha,
    rmsd,
)


class ModelWrapper(pl.LightningModule):
    def __init__(self, config):
        super(ModelWrapper, self).__init__()
        self.config = config
        self.model = AlphaFold(config)

        self.loss = AlphaFoldLoss(config.loss)

        self.ema = ExponentialMovingAverage(
            model=self.model, decay=config.ema.decay
        )
        
        self.cached_weights = None
        self.last_lr_step = -1

        self.aggregated_metrics = {}
        self.log_agg_every_n_steps = 50  # match Trainer(log_every_n_steps=50)

    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():
            # print("logging loss", loss_name, indiv_loss, flush=True)
            self.log(
                f"{phase}/{loss_name}", 
                indiv_loss, 
                on_step=train, on_epoch=(not train), logger=True, sync_dist=True
            )

            if train:
                agg_name = f"{phase}/{loss_name}_agg"
                if agg_name not in self.aggregated_metrics:
                    self.aggregated_metrics[agg_name] = []
                self.aggregated_metrics[agg_name].append(float(indiv_loss))
                self.log(
                    f"{phase}/{loss_name}_epoch",
                    indiv_loss,
                    on_step=False, on_epoch=True, logger=True, sync_dist=True
                )

        # print("logging validation metrics", flush=True)
        with torch.no_grad():
            other_metrics = self._compute_validation_metrics(
                batch, 
                outputs,
                superimposition_metrics=(not train)
            )

        for k, v in other_metrics.items():
            # print("logging metric", k, v, flush=True)
            if train:
                agg_name = f"{phase}/{k}_agg"
                if agg_name not in self.aggregated_metrics:
                    self.aggregated_metrics[agg_name] = []
                self.aggregated_metrics[agg_name].append(float(torch.mean(v)))
            self.log(
                f"{phase}/{k}",
                torch.mean(v),
                on_step=False, on_epoch=True, logger=True, sync_dist=True
            )

        if train and any([len(v) >= self.log_agg_every_n_steps for v in self.aggregated_metrics.values()]):
            for k, v in self.aggregated_metrics.items():
                print("logging agg", k, len(v), sum(v) / len(v), flush=True)
                self.log(k, sum(v) / len(v), on_step=True, on_epoch=False, logger=True, sync_dist=True)
                self.aggregated_metrics[k] = []

    def training_step(self, batch, batch_idx):
        if self.ema.device != batch["aatype"].device:
            self.ema.to(batch["aatype"].device)

        # ground_truth = batch.pop('gt_features', None)

        # Run the model
        # print("running model", round(time.time() % 10000, 3), flush=True)
        outputs = self(batch)

        # Remove the recycling dimension
        batch = tensor_tree_map(lambda t: t[..., -1], batch)

        # print("running loss", round(time.time() % 10000, 3), flush=True)
        # Compute loss
        loss, loss_breakdown = self.loss(
            outputs, batch, _return_breakdown=True
        )

        # Log it
        self._log(loss_breakdown, batch, outputs)
        # print("loss done", round(time.time() % 10000, 3), flush=True)


        return loss

    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:
            # 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"])

        # Run the model
        outputs = self(batch)
        batch = tensor_tree_map(lambda t: t[..., -1], batch)

        batch["use_clamped_fape"] = 0.

        # Compute loss and other metrics
        _, loss_breakdown = self.loss(
            outputs, batch, _return_breakdown=True
        )

        self._log(loss_breakdown, batch, outputs, train=False)
        
    def on_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 = {}
        
        all_gt_coords = batch["atom37_gt_positions"]
        all_pred_coords = outputs["final_atom_positions"]
        all_atom_mask = batch["atom37_atom_exists_in_gt"]

        rough_protein_atom_mask = torch.repeat_interleave(batch["protein_mask"], 37, dim=-1).view(*all_atom_mask.shape)
        protein_gt_coords = all_gt_coords * rough_protein_atom_mask[..., None]
        protein_pred_coords = all_pred_coords * rough_protein_atom_mask[..., None]
        protein_all_atom_mask = all_atom_mask * rough_protein_atom_mask

        rough_ligand_atom_mask = torch.repeat_interleave(batch["ligand_mask"], 37, dim=-1).view(*all_atom_mask.shape)
        ligand_gt_coords = all_gt_coords * rough_ligand_atom_mask[..., None]
        ligand_pred_coords = all_pred_coords * rough_ligand_atom_mask[..., None]
        ligand_all_atom_mask = all_atom_mask * rough_ligand_atom_mask

        # This is super janky for superimposition. Fix later
        protein_gt_coords_masked = protein_gt_coords * protein_all_atom_mask[..., None]
        protein_pred_coords_masked = protein_pred_coords * protein_all_atom_mask[..., None]
        ca_pos = residue_constants.atom_order["CA"]
        protein_gt_coords_masked_ca = protein_gt_coords_masked[..., ca_pos, :]
        protein_pred_coords_masked_ca = protein_pred_coords_masked[..., ca_pos, :]
        protein_atom_mask_ca = protein_all_atom_mask[..., ca_pos]

        ligand_gt_coords_single_atom = ligand_gt_coords[..., ca_pos, :]
        ligand_pred_coords_single_atom = ligand_pred_coords[..., ca_pos, :]
        ligand_gt_mask_single_atom = ligand_all_atom_mask[..., ca_pos]
    
        lddt_ca_score = lddt_ca(
            protein_pred_coords,
            protein_gt_coords,
            protein_all_atom_mask,
            eps=self.config.globals.eps,
            per_residue=False,
        )
   
        metrics["lddt_ca"] = lddt_ca_score
   
        drmsd_ca_score = drmsd(
            protein_pred_coords_masked_ca,
            protein_gt_coords_masked_ca,
            mask=protein_atom_mask_ca,  # still required here to compute n
        )
   
        metrics["drmsd_ca"] = drmsd_ca_score

        drmsd_intra_ligand_score = drmsd(
            ligand_pred_coords_single_atom,
            ligand_gt_coords_single_atom,
            mask=ligand_gt_mask_single_atom,
        )

        metrics["drmsd_intra_ligand"] = drmsd_intra_ligand_score

        # --- inter contacts
        gt_contacts = batch["gt_inter_contacts"]
        pred_contacts = torch.sigmoid(outputs["inter_contact_logits"].clone().detach()).squeeze(-1)
        pred_contacts = (pred_contacts > 0.5).float()
        pred_contacts = pred_contacts * batch["inter_pair_mask"]


        # Calculate True Positives, False Positives, and False Negatives
        tp = torch.sum((gt_contacts == 1) & (pred_contacts == 1))
        fp = torch.sum((gt_contacts == 0) & (pred_contacts == 1))
        fn = torch.sum((gt_contacts == 1) & (pred_contacts == 0))

        # Calculate Recall and Precision
        recall = tp / (tp + fn) if (tp + fn) > 0 else tp.float()
        precision = tp / (tp + fp) if (tp + fp) > 0 else tp.float()

        metrics["inter_contacts_recall"] = recall.clone().detach()
        metrics["inter_contacts_precision"] = precision.clone().detach()

        # print("inter_contacts recall", recall, "precision", precision, tp, fp, fn, torch.ones_like(gt_contacts).sum())

        # --- Affinity
        if True or batch["affinity_loss_factor"].sum() > 0.1:
            # print("affinity loss factor", batch["affinity_loss_factor"].sum())
            gt_affinity = batch["affinity"].squeeze(-1)
            affinity_linspace = torch.linspace(0, 15, 32, device=batch["affinity"].device)
            pred_affinity_1d = torch.sum(
                torch.softmax(outputs["affinity_1d_logits"].clone().detach(), -1) * affinity_linspace, dim=-1)

            pred_affinity_2d = torch.sum(
                torch.softmax(outputs["affinity_2d_logits"].clone().detach(), -1) * affinity_linspace, dim=-1)

            pred_affinity_cls = torch.sum(
                torch.softmax(outputs["affinity_cls_logits"].clone().detach(), -1) * affinity_linspace, dim=-1)

            aff_loss_factor = batch["affinity_loss_factor"].squeeze()

            metrics["affinity_dist_1d"] = (torch.abs(gt_affinity - pred_affinity_1d) * aff_loss_factor).sum() / aff_loss_factor.sum()
            metrics["affinity_dist_2d"] = (torch.abs(gt_affinity - pred_affinity_2d) * aff_loss_factor).sum() / aff_loss_factor.sum()
            metrics["affinity_dist_cls"] = (torch.abs(gt_affinity - pred_affinity_cls) * aff_loss_factor).sum() / aff_loss_factor.sum()
            metrics["affinity_dist_avg"] = (torch.abs(gt_affinity - (pred_affinity_cls + pred_affinity_1d + pred_affinity_2d) / 3) * aff_loss_factor).sum() / aff_loss_factor.sum()
            # print("affinity metrics", gt_affinity, pred_affinity_2d, aff_loss_factor, metrics["affinity_dist_1d"],
            #       metrics["affinity_dist_2d"], metrics["affinity_dist_cls"], metrics["affinity_dist_avg"])
        else:
            # print("skipping affinity metrics")
            pass
        if superimposition_metrics:
            superimposed_pred, alignment_rmsd, rots, transs = superimpose(
                protein_gt_coords_masked_ca, protein_pred_coords_masked_ca, protein_atom_mask_ca,
            )
            gdt_ts_score = gdt_ts(
                superimposed_pred, protein_gt_coords_masked_ca, protein_atom_mask_ca
            )
            gdt_ha_score = gdt_ha(
                superimposed_pred, protein_gt_coords_masked_ca, protein_atom_mask_ca
            )

            metrics["protein_alignment_rmsd"] = alignment_rmsd
            metrics["gdt_ts"] = gdt_ts_score
            metrics["gdt_ha"] = gdt_ha_score

            superimposed_ligand_coords = ligand_pred_coords_single_atom @ rots + transs[:, None, :]
            ligand_alignment_rmsds = rmsd(ligand_gt_coords_single_atom, superimposed_ligand_coords,
                                          mask=ligand_gt_mask_single_atom)
            metrics["ligand_alignment_rmsd"] = ligand_alignment_rmsds.mean()
            metrics["ligand_alignment_rmsd_under_2"] = torch.mean((ligand_alignment_rmsds < 2).float())
            metrics["ligand_alignment_rmsd_under_5"] = torch.mean((ligand_alignment_rmsds < 5).float())

            print("ligand rmsd:", ligand_alignment_rmsds)

        return metrics

    def configure_optimizers(self, 
        learning_rate: Optional[float] = None,
        eps: float = 1e-5,
    ) -> torch.optim.Adam:
        if learning_rate is None:
            learning_rate = self.config.globals.max_lr

        optimizer = torch.optim.Adam(
            self.model.parameters(), 
            lr=learning_rate, 
            eps=eps
        )

        if self.last_lr_step != -1:
            for group in optimizer.param_groups:
                if 'initial_lr' not in group:
                    group['initial_lr'] = learning_rate

        lr_scheduler = AlphaFoldLRScheduler(
            optimizer,
            last_epoch=self.last_lr_step,
            max_lr=self.config.globals.max_lr,
            start_decay_after_n_steps=10000,
            decay_every_n_steps=10000,
        )

        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": lr_scheduler,
                "interval": "step",
                "name": "AlphaFoldLRScheduler",
            }
        }

    def on_load_checkpoint(self, checkpoint):
        ema = checkpoint["ema"]
        self.ema.load_state_dict(ema)

    def on_save_checkpoint(self, checkpoint):
        checkpoint["ema"] = self.ema.state_dict()

    def resume_last_lr_step(self, lr_step):
        self.last_lr_step = lr_step


def override_config(base_config, overriding_config):
    for k, v in overriding_config.items():
        if isinstance(v, dict):
            base_config[k] = override_config(base_config[k], v)
        else:
            base_config[k] = v
    return base_config


def train(override_config_path: str):
    run_config = json.load(open(override_config_path, "r"))
    seed = 42
    seed_everything(seed, workers=True)
    output_dir = run_config["train_output_dir"]
    os.makedirs(output_dir, exist_ok=True)

    print("Starting train", time.time())
    config = model_config(
        run_config.get("stage", "initial_training"),
        train=True,
        low_prec=True
    )
    config = override_config(config, run_config.get("override_conf", {}))
    accumulate_grad_batches = run_config.get("accumulate_grad_batches", 1)
    print("config loaded", time.time())

    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    device_name = "cuda" if torch.cuda.is_available() else "cpu"
    # device_name = "mps" if device_name == "cpu" and torch.backends.mps.is_available() else device_name
    model_module = ModelWrapper(config)
    print("model loaded", time.time())

    # device_name = "cpu"

    # for debugging memory:
    # torch.cuda.memory._record_memory_history()

    if "train_input_dir" in run_config:
        data_module = OpenFoldDataModule(
            config=config.data,
            batch_seed=seed,
            train_data_dir=run_config["train_input_dir"],
            val_data_dir=run_config["val_input_dir"],
            train_epoch_len=run_config.get("train_epoch_len", 1000),
        )
    else:
        data_module = DockFormerDataModule(
            config=config.data,
            batch_seed=seed,
            train_data_file=run_config["train_input_file"],
            val_data_file=run_config["val_input_file"],
        )
    print("data module loaded", time.time())

    checkpoint_dir = os.path.join(output_dir, "checkpoint")
    ckpt_path = run_config.get("ckpt_path", get_latest_checkpoint(checkpoint_dir))

    if ckpt_path:
        print(f"Resuming from checkpoint: {ckpt_path}")
        sd = torch.load(ckpt_path)
        last_global_step = int(sd['global_step'])
        model_module.resume_last_lr_step(last_global_step)

    # Do we need this?
    data_module.prepare_data()
    data_module.setup("fit")

    callbacks = []

    mc = ModelCheckpoint(
        dirpath=checkpoint_dir,
        # every_n_epochs=1,
        every_n_train_steps=250,
        auto_insert_metric_name=False,
        save_top_k=1,
        save_on_train_epoch_end=True,  # before validation
    )

    mc2 = ModelCheckpoint(
        dirpath=checkpoint_dir,  # Directory to save checkpoints
        filename="step{step}_lig_rmsd{val/ligand_alignment_rmsd:.2f}",  # Filename format for best
        monitor="val/ligand_alignment_rmsd",  # Metric to monitor
        mode="min",  # We want the lowest `ligand_rmsd`
        save_top_k=1,  # Save only the best model based on `ligand_rmsd`
        every_n_epochs=1,  # Save a checkpoint every epoch
        auto_insert_metric_name=False
    )
    callbacks.append(mc)
    callbacks.append(mc2)

    lr_monitor = LearningRateMonitor(logging_interval="step")
    callbacks.append(lr_monitor)

    loggers = []

    wandb_project_name = "EvoDocker3"
    wandb_run_id_path = os.path.join(output_dir, "wandb_run_id.txt")

    # Initialize WandbLogger and save run_id
    local_rank = int(os.getenv('LOCAL_RANK', os.getenv("SLURM_PROCID", '0')))
    global_rank = int(os.getenv('GLOBAL_RANK', os.getenv("SLURM_LOCALID", '0')))
    print("ranks", os.getenv('LOCAL_RANK', 'd0'), os.getenv('local_rank', 'd0'), os.getenv('GLOBAL_RANK', 'd0'),
          os.getenv('global_rank', 'd0'), os.getenv("SLURM_PROCID", 'd0'), os.getenv('SLURM_LOCALID', 'd0'), flush=True)
    if local_rank == 0 and global_rank == 0 and not os.path.exists(wandb_run_id_path):
        wandb_logger = WandbLogger(project=wandb_project_name, save_dir=output_dir)
        with open(wandb_run_id_path, 'w') as f:
            f.write(wandb_logger.experiment.id)
        wandb_logger.experiment.config.update(run_config, allow_val_change=True)
    else:
        # Necessary for multi-node training https://github.com/rstrudel/segmenter/issues/22
        while not os.path.exists(wandb_run_id_path):
            print(f"Waiting for run_id file to be created ({local_rank})", flush=True)
            time.sleep(1)
        with open(wandb_run_id_path, 'r') as f:
            run_id = f.read().strip()
        wandb_logger = WandbLogger(project=wandb_project_name, save_dir=output_dir, resume='must', id=run_id)
    loggers.append(wandb_logger)

    strategy_params = {"strategy": "auto"}
    if run_config.get("multi_node", False):
        strategy_params["strategy"] = "ddp"
        # strategy_params["strategy"] = "ddp_find_unused_parameters_true" # this causes issues with checkpointing...
        strategy_params["num_nodes"] = run_config["multi_node"]["num_nodes"]
        strategy_params["devices"] = run_config["multi_node"]["devices"]

    trainer = pl.Trainer(
        accelerator=device_name,
        default_root_dir=output_dir,
        **strategy_params,
        reload_dataloaders_every_n_epochs=1,
        accumulate_grad_batches=accumulate_grad_batches,
        check_val_every_n_epoch=run_config.get("check_val_every_n_epoch", 10),
        callbacks=callbacks,
        logger=loggers,
        # profiler=AdvancedProfiler(),
    )

    print("Starting fit", time.time())
    trainer.fit(
        model_module,
        datamodule=data_module,
        ckpt_path=ckpt_path,
    )

    # profiler_results = trainer.profiler.summary()
    # print(profiler_results)

    # torch.cuda.memory._dump_snapshot("my_train_snapshot.pickle")
    # view on https://pytorch.org/memory_viz


if __name__ == "__main__":
    if len(sys.argv) > 1:
        train(sys.argv[1])
    else:
        train(os.path.join(os.path.dirname(__file__), "run_config.json"))