model/barlow_twins.py

import torch
torch.manual_seed(42)
torch.backends.cudnn.deterministic = True
from torch import nn
import numpy as np
from typing import *
from datetime import datetime
import os
import pickle
import inspect
from tqdm.auto import trange
import spaces

from model.base_model import BaseModel


class BarlowTwins(BaseModel):
    def __init__(
            self,
            n_bits: int = 1024,
            aa_emb_size: int = 1024,
            enc_n_neurons: int = 512,
            enc_n_layers: int = 2,
            proj_n_neurons: int = 2048,
            proj_n_layers: int = 2,
            embedding_dim: int = 512,
            act_function: str = "relu",
            loss_weight: float = 0.005,
            batch_size: int = 512,
            optimizer: str = "adamw",
            momentum: float = 0.9,
            learning_rate: float = 0.0001,
            betas: tuple = (0.9, 0.999),
            weight_decay: float = 1e-3,
            step_size: int = 10,
            gamma: float = 0.1,
            verbose: bool = True,
    ):
        super().__init__()
        
        self.enc_aa = None
        self.enc_mol = None
        self.proj = None
        
        self.scheduler = None
        self.optimizer = None

        # store input in dict
        self.param_dict = {
            "act_function": self.activation_dict[
                act_function
            ],  # which activation function to use among dict options
            "loss_weight": loss_weight,  # off-diagonal cross correlation loss weight
            "batch_size": batch_size,  # samples per gradient step
            "learning_rate": learning_rate,  # update step magnitude when training
            "betas": betas,  # momentum hyperparameter for adam-like optimizers
            "step_size": step_size,  # decay period for the learning rate
            "gamma": gamma,  # decay coefficient for the learning rate
            "optimizer": self.optimizer_dict[
                optimizer
            ],  # which optimizer to use among dict options
            "momentum": momentum,  # momentum hyperparameter for SGD
            "enc_n_neurons": enc_n_neurons,  # neurons to use for the mlp encoder
            "enc_n_layers": enc_n_layers,  # number of hidden layers in the mlp encoder
            "proj_n_neurons": proj_n_neurons,  # neurons to use for the mlp projector
            "proj_n_layers": proj_n_layers,  # number of hidden layers in the mlp projector
            "embedding_dim": embedding_dim,  # latent space dim for downstream tasks
            "weight_decay": weight_decay,  # l2 regularization for linear layers
            "verbose": verbose,  # whether to print feedback
            "radius": "Not defined yet",  # fingerprint radius
            "n_bits": n_bits,  # fingerprint bit size
            "aa_emb_size": aa_emb_size,  # aa embedding size
        }

        # create history dictionary
        self.history = {
            "train_loss": [],
            "on_diag_loss": [],
            "off_diag_loss": [],
            "validation_loss": [],
        }

        # run NN architecture construction method
        self.construct_model()

        # run scheduler construction method
        self.construct_scheduler()

        # print if necessary
        if self.param_dict["verbose"] is True:
            self.print_config()

    @staticmethod
    def __validate_inputs(locals_dict) -> None:
        # get signature types from __init__
        init_signature = inspect.signature(BarlowTwins.__init__)

        # loop over all chosen arguments
        for param_name, param_value in locals_dict.items():
            # skip self
            if param_name != "self":
                # check that parameter exists
                if param_name in init_signature.parameters:
                    # check that param is correct type
                    expected_type = init_signature.parameters[param_name].annotation
                    assert isinstance(
                        param_value, expected_type
                    ), f"[BT]: Type mismatch for parameter '{param_name}'"
                else:
                    raise ValueError(f"[BT]: Unexpected parameter '{param_name}'")

    def construct_mlp(self, input_units, layer_units, n_layers, output_units) -> nn.Sequential:

        # make empty list to fill
        mlp_list = []

        # make lists defining layer sizes (input + n_neurons*n_layers + embedding_dim)
        units = [input_units] + [layer_units] * n_layers

        # add layer stack (linear -> batchnorm -> dropout -> activation)
        for i in range(len(units) - 1):
            mlp_list.append(nn.Linear(units[i], units[i + 1]))
            mlp_list.append(nn.BatchNorm1d(units[i + 1]))
            mlp_list.append(self.param_dict["act_function"]())

        # add final linear layer
        mlp_list.append(nn.Linear(units[-1], output_units))

        return nn.Sequential(*mlp_list)

    def construct_model(self) -> None:
        # create fingerprint transformer
        self.enc_mol = self.construct_mlp(
            self.param_dict["n_bits"],
            self.param_dict["enc_n_neurons"],
            self.param_dict["enc_n_layers"],
            self.param_dict["embedding_dim"],
        )

        # create aa transformer
        self.enc_aa = self.construct_mlp(
            self.param_dict["aa_emb_size"],
            self.param_dict["enc_n_neurons"],
            self.param_dict["enc_n_layers"],
            self.param_dict["embedding_dim"],
        )

        # create mlp projector
        self.proj = self.construct_mlp(
            self.param_dict["embedding_dim"],
            self.param_dict["proj_n_neurons"],
            self.param_dict["proj_n_layers"],
            self.param_dict["proj_n_neurons"],
        )

        # print if necessary
        if self.param_dict["verbose"] is True:
            print("[BT]: Model constructed successfully")

    def construct_scheduler(self):
        # make optimizer
        self.optimizer = self.param_dict["optimizer"](
            list(self.enc_mol.parameters())
            + list(self.enc_aa.parameters())
            + list(self.proj.parameters()),
            lr=self.param_dict["learning_rate"],
            betas=self.param_dict["betas"],
            # momentum=self.param_dict["momentum"],
            weight_decay=self.param_dict["weight_decay"],
        )

        # wrap optimizer in scheduler
        """
         self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            self.optimizer,
            T_max=self.param_dict["step_size"], # T_0
            # eta_min=1e-7,
            verbose=True
        ) 

        self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer,
            patience=self.param_dict["step_size"],
            verbose=True
        )
        """
        self.scheduler = torch.optim.lr_scheduler.StepLR(
            self.optimizer,
            step_size=self.param_dict["step_size"],
            gamma=self.param_dict["gamma"],
        )

        # print if necessary
        if self.param_dict["verbose"] is True:
            print("[BT]: Optimizer constructed successfully")

    def switch_mode(self, is_training: bool):
        if is_training:
            self.enc_mol.train()
            self.enc_aa.train()
            self.proj.train()
        else:
            self.enc_mol.eval()
            self.enc_aa.eval()
            self.proj.eval()

    @staticmethod
    def normalize_projection(tensor: torch.tensor) -> torch.tensor:
        means = torch.mean(tensor, axis=0)
        std = torch.std(tensor, axis=0)
        centered = torch.add(tensor, -means)
        scaled = torch.div(centered, std)

        return scaled

    def compute_loss(
        self,
        mol_embedding: torch.tensor,
        aa_embedding: torch.tensor,
    ) -> torch.tensor:

        # empirical cross-correlation matrix
        mol_embedding = self.normalize_projection(mol_embedding).T
        aa_embedding = self.normalize_projection(aa_embedding)
        c = mol_embedding @ aa_embedding

        # normalize by number of samples
        c.div_(self.param_dict["batch_size"])

        # compute elements on diagonal
        on_diag = torch.diagonal(c).add_(-1).pow_(2).sum()

        # compute elements off diagonal
        n, m = c.shape
        off_diag = c.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()
        off_diag = off_diag.pow_(2).sum() * self.param_dict["loss_weight"]

        return on_diag, off_diag

    def forward(
        self, mol_data: torch.tensor, aa_data: torch.tensor, is_training: bool = True
    ) -> torch.tensor:

        # switch according to input
        self.switch_mode(is_training)

        # get embeddings
        mol_embeddings = self.enc_mol(mol_data)
        aa_embeddings = self.enc_aa(aa_data)

        # get projections
        mol_proj = self.proj(mol_embeddings)
        aa_proj = self.proj(aa_embeddings)

        # compute loss
        on_diag, off_diag = self.compute_loss(mol_proj, aa_proj)

        return on_diag, off_diag

    def train(
        self,
        train_data: torch.utils.data.DataLoader,
        val_data: torch.utils.data.DataLoader = None,
        num_epochs: int = 20,
        patience: int = None,
    ):
        if self.param_dict["verbose"] is True:
            print("[BT]: Training started")

        if patience is None:
            patience = 2 * self.param_dict["step_size"]

        pbar = trange(num_epochs, desc="[BT]: Epochs", leave=False, colour="blue")

        for epoch in pbar:
            # initialize loss containers
            train_loss = 0.0
            on_diag_loss = 0.0
            off_diag_loss = 0.0
            val_loss = 0.0

            # loop over training set
            for _, (mol_data, aa_data) in enumerate(train_data):
                # reset grad
                self.optimizer.zero_grad()

                # compute train loss for batch
                on_diag, off_diag = self.forward(mol_data, aa_data, is_training=True)
                t_loss = on_diag + off_diag

                # backpropagation and optimization
                t_loss.backward()
                """
                nn.utils.clip_grad_norm_(
                    list(self.enc_mol.parameters()) +
                    list(self.enc_aa.parameters()) +
                    list(self.proj.parameters()),
                    1
                )
                """
                self.optimizer.step()

                # add i-th loss to training container
                train_loss += t_loss.item()
                on_diag_loss += on_diag.item()
                off_diag_loss += off_diag.item()

            # add mean epoch loss for train data to history dictionary
            self.history["train_loss"].append(train_loss / len(train_data))
            self.history["on_diag_loss"].append(on_diag_loss / len(train_data))
            self.history["off_diag_loss"].append(off_diag_loss / len(train_data))

            # define msg to be printed
            msg = (
                f"[BT]: Epoch [{epoch + 1}/{num_epochs}], "
                f"Train loss: {train_loss / len(train_data):.3f}, "
                f"On diagonal: {on_diag_loss / len(train_data):.3f}, "
                f"Off diagonal: {off_diag_loss / len(train_data):.3f} "
            )

            # loop over validation set (if present)
            if val_data is not None:

                for _, (mol_data, aa_data) in enumerate(val_data):
                    # compute val loss for batch
                    on_diag_v_loss, off_diag_v_loss = self.forward(
                        mol_data, aa_data, is_training=False
                    )

                    # add i-th loss to val container
                    v_loss = on_diag_v_loss + off_diag_v_loss
                    val_loss += v_loss.item()

                # add mean epoc loss for val data to history dictionary
                self.history["validation_loss"].append(val_loss / len(val_data))

                # add val loss to msg
                msg += f", Val loss: {val_loss / len(val_data):.3f}"

                # early stopping
                if self.early_stopping(patience=patience):
                    break

                pbar.set_postfix(
                    {
                        "train loss": train_loss / len(train_data),
                        "val loss": val_loss / len(val_data),
                    }
                )

            else:
                pbar.set_postfix({"train loss": train_loss / len(train_data)})

            # update scheduler
            self.scheduler.step()  # val_loss / len(val_data)

            if self.param_dict["verbose"] is True:
                print(msg)

        if self.param_dict["verbose"] is True:
            print("[BT]: Training finished")

    @spaces.GPU
    def encode(
        self, vector: np.ndarray, mode: str = "embedding", normalize: bool = True, encoder: str = "mol"
    ) -> np.ndarray:
        """
        Encodes a given vector using the Barlow Twins model.

        Args:
        - vector (np.ndarray): the input vector to encode
        - mode (str): the mode to use for encoding, either "embedding" or "projection"
        - normalize (bool): whether to L2 normalize the output vector

        Returns:
        - np.ndarray: the encoded vector
        """

        # set mol encoder to eval mode
        self.switch_mode(is_training=False)

        # convert from numpy to tensor
        if type(vector) is not torch.Tensor:
            vector = torch.from_numpy(vector)

        # if oly one molecule pair is passed, add a batch dimension
        if len(vector.shape) == 1:
            vector = vector.unsqueeze(0)

        # get representation
        if encoder == "mol":
            embedding = self.enc_mol(vector)
            if mode == "projection":
                embedding = self.proj(embedding)
        elif encoder == "aa":
            embedding = self.enc_aa(vector)
            if mode == "projection":
                embedding = self.proj(embedding)
        else:
            raise ValueError("[BT]: Encoder not recognized")

        # L2 normalize (optional)
        if normalize:
            embedding = torch.nn.functional.normalize(embedding)

        # convert back to numpy
        return embedding.cpu().detach().numpy()

    @spaces.GPU
    def zero_shot(
        self, mol_vector: np.ndarray, aa_vector: np.ndarray, l2_norm: bool = True, device: str = "cpu"
    ) -> np.ndarray:

        # disable training
        self.switch_mode(is_training=False)

        # cast aa vectors (pos and neg) to correct size, force single precision
        # to both
        mol_vector = np.array(mol_vector, dtype=np.float32)
        aa_vector = np.array(aa_vector, dtype=np.float32)

        # convert to tensors
        mol_vector = torch.from_numpy(mol_vector).to(device)
        aa_vector = torch.from_numpy(aa_vector).to(device)

        # get embeddings
        mol_embedding = self.encode(mol_vector, normalize=l2_norm, encoder="mol")
        aa_embedding = self.encode(aa_vector, normalize=l2_norm, encoder="aa")

        # concat mol and aa embeddings
        concat = np.concatenate((mol_embedding, aa_embedding), axis=1)
        return concat
    
    def zero_shot_explain(
        self, mol_vector, aa_vector, l2_norm: bool = True, device: str = "cpu"
    ):
        self.switch_mode(is_training=False)

        mol_embedding = self.encode(mol_vector, normalize=l2_norm, encoder="mol")
        aa_embedding = self.encode(aa_vector, normalize=l2_norm, encoder="aa")

        return torch.cat((mol_embedding, aa_embedding), dim=1)

    def consume_preprocessor(self, preprocessor) -> None:
        # save attributes related to fingerprint generation from
        # preprocessor object
        self.param_dict["radius"] = preprocessor.radius
        self.param_dict["n_bits"] = preprocessor.n_bits

    def save_model(self, path: str) -> None:
        # get current date and time for the filename
        now = datetime.now()
        formatted_date = now.strftime("%d%m%Y")
        formatted_time = now.strftime("%H%M")
        folder_name = f"{formatted_date}_{formatted_time}"

        # make full path string and folder
        folder_path = path + "/" + folder_name
        os.makedirs(folder_path)

        # make paths for weights, config and history
        weight_path = folder_path + "/weights.pt"
        param_path = folder_path + "/params.pkl"
        history_path = folder_path + "/history.json"

        # save each Sequential state dict in one object to the path
        torch.save(
            {
                "enc_mol": self.enc_mol.state_dict(),
                "enc_aa": self.enc_aa.state_dict(),
                "proj": self.proj.state_dict(),
            },
            weight_path,
        )

        # dump params in pkl
        with open(param_path, "wb") as file:
            pickle.dump(self.param_dict, file)

        # dump history in json
        with open(history_path, "wb") as file:
            pickle.dump(self.history, file)

        # print if verbose is True
        if self.param_dict["verbose"] is True:
            print(f"[BT]: Model saved at {folder_path}")

    @spaces.GPU
    def load_model(self, path: str) -> None:
        # make weights, config and history paths
        weights_path = path + "/weights.pt"
        param_path = path + "/params.pkl"
        history_path = path + "/history.json"

        # load weights, history and params
        checkpoint = torch.load(weights_path, map_location=self.device)
        with open(param_path, "rb") as file:
            param_dict = pickle.load(file)
        with open(history_path, "rb") as file:
            history = pickle.load(file)

        # construct model again, overriding old verbose key with new instance
        verbose = self.param_dict["verbose"]
        self.param_dict = param_dict
        self.param_dict["verbose"] = verbose
        self.history = history
        self.construct_model()

        # set weights in Sequential models
        self.enc_mol.load_state_dict(checkpoint["enc_mol"])
        self.enc_aa.load_state_dict(checkpoint["enc_aa"])
        self.proj.load_state_dict(checkpoint["proj"])

        # recreate scheduler and optimizer in order to add new weights
        # to graph
        self.construct_scheduler()

        # print if verbose is True
        if self.param_dict["verbose"] is True:
            print(f"[BT]: Model loaded from {path}")
            print("[BT]: Loaded parameters:")
            print(self.param_dict)
    
    @spaces.GPU
    def move_to_device(self, device) -> None:
        # move each Sequential model to device
        self.enc_mol.to(device)
        self.enc_aa.to(device)
        self.proj.to(device)
        self.device = device