Training.py

import os
import glob
import torch
import lightning as pl
from lightning.pytorch.loggers import WandbLogger
from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint, LearningRateMonitor
import torch.nn as nn
import numpy as np
from pandas import read_csv
from torch import  Tensor
import torch.nn.functional as F
import re
import seaborn as sns
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torch.utils.data import Dataset

import wandb
from torchmetrics import MetricCollection
from torchmetrics.classification import BinaryAccuracy, BinaryRecall, BinaryPrecision, BinaryConfusionMatrix, BinaryF1Score, MultilabelAccuracy, MultilabelF1Score, MultilabelConfusionMatrix,MultilabelPrecision,MultilabelRecall

from torchmetrics.regression import MeanSquaredError,R2Score,MeanAbsoluteError

from src.model.Model import TransformerMintomics
from argparse import ArgumentParser
import scipy.signal as signal
from src.dataset.PrepareDataset import Psedu_data, Data2target, gene2protein
from src.dataset.prepare_test_data import Data2target_test


AVAIL_GPUS = [1,2]
NUM_NODES = 1
BATCH_SIZE = 1
DATALOADERS = 1
ACCELERATOR = "gpu"
EPOCHS = 1
ATT_HEAD = 1
ENCODE_LAYERS = 1
DATASET_DIR = "./"

#label_dict = read_csv(DATASET_DIR+"/Dataset/Labels_proc/Labels_control.csv",index_col=0)
#label_dict = read_csv(DATASET_DIR+"/Dataset/Labels_proc/Labels_control.csv",index_col=0)
#Num_classes = len(label_dict)
Num_classes = 3060
"""

torch.set_default_tensor_type(torch.FloatTensor)  # Ensure that the default tensor type is FloatTensor
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # Choose the device you want to use

if device.type == "cuda":
    torch.backends.cudnn.benchmark = True  # Enable cuDNN auto-tuner to find the best algorithm to use for hardware
    torch.set_default_tensor_type(torch.cuda.FloatTensor)  # Set the default tensor type to CUDA FloatTensor
    torch.set_float32_matmul_precision('medium')  # Set Tensor Core precision to medium

"""

CHECKPOINT_PATH = f"{DATASET_DIR}/Trainings/tempo"
os.makedirs(CHECKPOINT_PATH, exist_ok=True)


class Mintomics(pl.LightningModule):
    def __init__(self, learning_rate=1e-4,attn_head=ATT_HEAD,encoder_layers=ENCODE_LAYERS,n_class=1, num_labels=76, **model_kwargs):
        super().__init__()
        
        self.save_hyperparameters()
        self.model = TransformerMintomics(attn_head=attn_head,encoder_layers=encoder_layers,n_class=n_class,**model_kwargs)
        self.loss_fn = nn.BCEWithLogitsLoss()
       
        self.metrics_class = MetricCollection([MultilabelAccuracy(num_labels=num_labels,average='micro'),
                                              MultilabelPrecision(num_labels=num_labels,average='micro'),
                                              MultilabelF1Score(num_labels=num_labels,average='micro'),
                                              MultilabelRecall(num_labels=num_labels,average='micro')])
        #self.metrics_class = MetricCollection([BinaryAccuracy(),
        #                                 BinaryPrecision(),
        #                                 BinaryRecall(),
        #                                 BinaryF1Score()])

        self.train_metrics_class = self.metrics_class.clone(prefix="train_")
  
        self.valid_metrics_class = self.metrics_class.clone(prefix="valid_")
  
        self.test_metrics_class = self.metrics_class.clone(prefix="test_")
  

    def forward(self, pest_sample):
        x = self.model(pest_sample)
        return x
    
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
        lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer=optimizer, mode='min', factor=0.1, patience=20, eps=1e-10)
        metric_to_track = 'valid_loss'
        return{'optimizer':optimizer,
               'lr_scheduler':lr_scheduler,
               'monitor':metric_to_track}
    
    def training_step(self,batch,batch_idx):
        batch_data = batch[0]
        
        
        inf = batch[2]
        #print(batch_data.shape)
        y_hat,_ = self.forward(batch_data)
        ##print(y_hat.shape)
        batch_label_class = batch[3].to(self.device)
        
        class_pred = y_hat[:, inf[0, 1]]
        
        #target = torch.transpose(batch_label_class, 1, 2)
        target = batch_label_class[:, inf[0, 1]]
       
        #batch_label_class = batch_label_class[:,None].cuda()

        #class_pred = y_hat.view(-1) 
        
        loss_class = self.loss_fn(class_pred,target.float())
        metric_log_class = self.train_metrics_class(class_pred, target)
        self.log_dict(metric_log_class, on_step=True, on_epoch=True)
        loss = (loss_class)
        self.log('train_loss',loss, on_step=True, on_epoch=True, sync_dist=True)
        return loss
    
    def validation_step(self,batch,batch_idx):
        batch_data = batch[0]
        
        
        inf = batch[2]
        #print(batch_data.shape)
        y_hat,_ = self.forward(batch_data)
        ##print(y_hat.shape)
        batch_label_class = batch[3].to(self.device)
        
        class_pred = y_hat[:, inf[0, 1]]
        
        #target = torch.transpose(batch_label_class, 1, 2)
        target = batch_label_class[:, inf[0, 1]]
       
        #batch_label_class = batch_label_class[:,None].cuda()

        #class_pred = y_hat.view(-1) 
        
        loss_class = self.loss_fn(class_pred,target.float())
        metric_log_class = self.valid_metrics_class(class_pred, target)
        self.log_dict(metric_log_class, on_step=True, on_epoch=True)
        loss = (loss_class)
        self.log('valid_loss',loss, on_step=True, on_epoch=True, sync_dist=True)
        return loss
       
    
    def test_step(self,batch, batch_idx):
        batch_data = batch[0]
        inf = batch[2]
        y_hat,attnt = self.forward(batch_data)
        batch_label_class = batch[3].to(self.device)
        
        class_pred = y_hat[:, inf[0, 1]]
        
        #target = torch.transpose(batch_label_class, 1, 2)
        target = batch_label_class[:, inf[0, 1]]
       
        #batch_label_class = batch_label_class[:,None].cuda()

        #class_pred = y_hat.view(-1) 
        
        loss_class = self.loss_fn(class_pred,target.float())
        metric_log_class = self.test_metrics_class(class_pred, target)
        self.log_dict(metric_log_class)
        loss = (loss_class)
        self.log('test_loss',loss, on_step=True, on_epoch=True, sync_dist=True)
       
        #conf_mat = BinaryConfusionMatrix().to("cuda")
        #conf_vals = conf_mat(class_pred, batch_label_class.squeeze())
        #print("Test Data Confusion Matrix: \n")
        #print(conf_vals)
        return {f'preds_class' : class_pred, f'targets_class' :target,f'attention':attnt,f'inf':inf}
        
    def test_epoch_end(self, outputs):
        # Log individual results for each dataset
        
        #for i  in range(len(outputs)):
            dataset_outputs = outputs
            
            #torch.save(dataset_outputs,"Predictions.pt")
            class_preds = torch.cat([x[f'preds_class'] for x in dataset_outputs])
            class_targets = torch.cat([x[f'targets_class'] for x in dataset_outputs])
            conf_mat = BinaryConfusionMatrix().to(self.device)
            conf_vals = conf_mat(class_preds, class_targets)
            fig = sns.heatmap(conf_vals.cpu() , annot=True, cmap="Blues", fmt="d")
            ind = torch.nonzero(class_targets[0,:]>0.5)
            attention = torch.cat([x[f'attention'] for x in dataset_outputs]).squeeze()
            inf = torch.cat([x[f'inf'] for x in dataset_outputs]).squeeze()
            attention1 = attention[:,inf[0,:]]
            attention2 = attention1[:,ind].squeeze()
            print(inf.shape, attention2.shape)
            # Get top 100 genes along rows for all columns
            top_genes_values, top_genes_indices = torch.topk(attention2, k=20, dim=0)
            mask = torch.zeros_like(attention2)
            
            mask[top_genes_indices, torch.arange(attention2.shape[1])] = 1.0
            print(mask)
            # Multiply the mask with the selected portion to keep only the top genes values
            attention2 = attention2 * mask
            # Calculate the hierarchical clustering
            # Calculate the hierarchical clustering
            #row_linkage = hierarchy.linkage(attention2, method='average')
            #col_linkage = hierarchy.linkage(attention2.T, method='average')

            # Reorder the matrix rows and columns based on the clustering
            #idx_row = hierarchy.dendrogram(row_linkage, no_plot=True)['leaves']
            #idx_col = hierarchy.dendrogram(col_linkage, no_plot=True)['leaves']
            # Calculate the hierarchical clustering
            #row_clusters = fastcluster.linkage(attention2, method='average')
            #col_clusters = fastcluster.linkage(attention2.T, method='average')
            
            # Plot the dendrogram for rows
            fig1 = plt.figure(figsize=(10, 20))
            sns.heatmap(attention2, cmap='rocket_r')
            plt.show()
            # Plot the reordered matrix
            #fig1 = plt.figure(figsize=(10, 10))
            #sns.clustermap(attention2, cmap='bone')
            #plt.show()
            #attention = self.model.encod.self_attn.
            #fig1 = plt.figure(figsize=(50, 100))
            #ax = fig1.add_subplot(111)
            #cax = ax.matshow(attention2.cpu().numpy(), cmap='bone')
            #cax.autoscale()
            #fig1.colorbar(cax)
            
            wandb.log({f"conf_mat" : wandb.Image(fig),"attentions":wandb.Image(fig1)})
            
            return super().test_epoch_end(outputs)

    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = ArgumentParser(parents=[parent_parser], add_help=False)
        parser.add_argument('--learning_rate', type=float, default=1e-4)
        parser.add_argument('--attn_head',type=int,default=ATT_HEAD)
        parser.add_argument('--encoder_layers',type=int,default=ENCODE_LAYERS)
        parser.add_argument('--n_class',type=int,default=1)
        return parser


def _logged_series(frame, column):
    if column not in frame.columns:
        return None
    series = frame.dropna(subset=[column])
    if series.empty:
        return None
    return series


def _smooth(values):
    window = max(1, len(values) // 50)
    return values.rolling(window, min_periods=1, center=True).mean()


def _style_axes(axes):
    for ax in axes:
        ax.tick_params(labelsize=16)
        ax.xaxis.label.set_size(18)
        ax.yaxis.label.set_size(18)
        ax.title.set_size(20)


def save_training_curves(metrics_csv, loss_png, metrics_png):
    frame = read_csv(metrics_csv)
    fig, axes = plt.subplots(1, 2, figsize=(14, 5.5), sharey=True)
    train_loss = _logged_series(frame, 'train_loss_step')
    valid_loss = _logged_series(frame, 'valid_loss_step')
    if train_loss is not None:
        axes[0].plot(train_loss['step'], _smooth(train_loss['train_loss_step']))
    if valid_loss is not None:
        axes[1].plot(valid_loss['step'], _smooth(valid_loss['valid_loss_step']), color='C1')
    axes[0].set_title('training loss')
    axes[1].set_title('validation loss')
    for ax in axes:
        ax.set_xlabel('step')
        ax.set_ylabel('loss')
    _style_axes(axes)
    fig.tight_layout()
    fig.savefig(loss_png, dpi=360)
    plt.close(fig)

    metric_names = [
        ('MultilabelAccuracy', 'accuracy'),
        ('MultilabelPrecision', 'precision'),
        ('MultilabelRecall', 'recall'),
        ('MultilabelF1Score', 'F1'),
    ]
    fig, axes = plt.subplots(1, 2, figsize=(14, 5.5), sharey=True)
    for key, label in metric_names:
        train_col = f'train_{key}_step' if f'train_{key}_step' in frame.columns else f'train_{key}'
        valid_col = f'valid_{key}_step' if f'valid_{key}_step' in frame.columns else f'valid_{key}'
        train_metric = _logged_series(frame, train_col)
        valid_metric = _logged_series(frame, valid_col)
        if train_metric is not None:
            axes[0].plot(train_metric['step'], _smooth(train_metric[train_col]), label=label)
        if valid_metric is not None:
            axes[1].plot(valid_metric['step'], _smooth(valid_metric[valid_col]), label=label)
    axes[0].set_title('training metrics')
    axes[1].set_title('validation metrics')
    for ax in axes:
        ax.set_xlabel('step')
        ax.set_ylabel('score')
        ax.set_ylim(0, 1)
        ax.legend(fontsize=16)
    _style_axes(axes)
    fig.tight_layout()
    fig.savefig(metrics_png, dpi=360)
    plt.close(fig)


def save_heldout_confusion(model, dataset, out_png, device):
    model.eval()
    model.to(device)
    sample = dataset[0]
    features = sample[0].unsqueeze(0).to(device)
    protein_index = sample[2][1].long().to(device)
    labels = sample[3].to(device).reshape(-1)
    with torch.no_grad():
        logits, _ = model(features)
    class_preds = logits[0, protein_index]
    class_targets = labels[protein_index]
    confusion = BinaryConfusionMatrix().to(device)
    values = confusion(class_preds, class_targets.int()).detach().cpu().numpy()
    fig, ax = plt.subplots(figsize=(6, 5))
    sns.heatmap(values, annot=True, cmap='Blues', fmt='d', ax=ax, annot_kws={'size': 18})
    ax.tick_params(labelsize=16)
    colorbar = ax.collections[0].colorbar if ax.collections else None
    if colorbar is not None:
        colorbar.ax.tick_params(labelsize=16)
    fig.tight_layout()
    fig.savefig(out_png, dpi=360)
    plt.close(fig)
    return values


def train_mintomics_classifier():
    pl.seed_everything(42)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    parser = ArgumentParser()
    parser = Mintomics.add_model_specific_args(parser)
    parser.add_argument('--num_gpus', type=int, default=1,
                        help="Number of GPUs to use (e.g. -1 = all available GPUs)")
    parser.add_argument('--nodes', type=int, default=NUM_NODES, help="Number of nodes to use")
    parser.add_argument('--num_epochs', type=int, default=EPOCHS, help="Number of epochs")
    parser.add_argument('--batch_size', default=BATCH_SIZE, type=int,
                        help="effective_batch_size = batch_size * num_gpus * num_nodes")
    parser.add_argument('--num_dataloader_workers', type=int, default=DATALOADERS)
    parser.add_argument('--entity_name', type=str, default='aghktb', help="Weights and Biases entity name")
    parser.add_argument('--project_name', type=str, default='Mintomics',
                        help="Weights and Biases project name")
    parser.add_argument('--save_dir', type=str, default='retrain', help="Directory in which to save models")
    parser.add_argument('--abundance_threshold', type=float, default=0.8,
                        help="Protein abundance cutoff for the positive class")
    parser.add_argument('--unit_test', type=int, default=0,
                        help="helps in debug, this touches all the parts of code."
                             "Enter True or num of batch you want to send, " "eg. 1 or 7")
    args = parser.parse_args()

    save_PATH = os.path.join(DATASET_DIR, "Trainings", args.save_dir)
    os.makedirs(save_PATH, exist_ok=True)

    dataset_train = Data2target(stage='train', size = 4000, pertage = 0.15, abundance_threshold=args.abundance_threshold)
    dataset_valid = Data2target(stage='valid', size = 1000, pertage = 0.15, abundance_threshold=args.abundance_threshold)
    dataset_test = Data2target_test(stage='test', size=1, pertage=0.0, abundance_threshold=args.abundance_threshold)
    train_loader = DataLoader(dataset=dataset_train, batch_size=args.batch_size, shuffle=True, num_workers=args.num_dataloader_workers)
    valid_loader = DataLoader(dataset=dataset_valid, batch_size=args.batch_size, shuffle=False, num_workers=args.num_dataloader_workers)

    num_labels = int(dataset_train.info.shape[-1])
    model = Mintomics(learning_rate=1e-4, n_class=Num_classes, num_labels=num_labels)
    checkpoint_callback = ModelCheckpoint(monitor='valid_loss', save_top_k=10, dirpath=save_PATH, filename='mintomics_{epoch:02d}_{valid_loss:.6f}')
    lr_monitor = LearningRateMonitor(logging_interval='epoch')
    early_stopping_callback = EarlyStopping(monitor='valid_loss', mode='min', min_delta=0.0, patience=30)
    from lightning.pytorch.loggers import CSVLogger
    csv_logger = CSVLogger(save_dir=save_PATH, name='metrics')
    accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'
    devices = args.num_gpus if accelerator == 'gpu' else 1
    trainer = pl.Trainer(
        accelerator=accelerator,
        devices=devices,
        num_nodes=args.nodes,
        max_epochs=args.num_epochs,
        log_every_n_steps=1,
        enable_model_summary=True,
        callbacks=[checkpoint_callback, lr_monitor, early_stopping_callback],
        logger=csv_logger,
        deterministic=False,
        limit_train_batches=args.unit_test if args.unit_test else 1.0,
        limit_val_batches=args.unit_test if args.unit_test else 1.0,
    )
    trainer.fit(model, train_loader, valid_loader)
    metrics_csv = os.path.join(csv_logger.log_dir, 'metrics.csv')
    save_training_curves(
        metrics_csv,
        os.path.join(save_PATH, 'train_val_loss_curves.png'),
        os.path.join(save_PATH, 'train_val_metrics_curves.png'),
    )
    best_path = trainer.checkpoint_callback.best_model_path
    best_model = Mintomics.load_from_checkpoint(best_path)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    values = save_heldout_confusion(best_model, dataset_test, os.path.join(save_PATH, 'heldout_confusion_matrix.png'), device)
    print('confusion_matrix', values.tolist())
    print('loss_curves', os.path.join(save_PATH, 'train_val_loss_curves.png'))
    print('metrics_curves', os.path.join(save_PATH, 'train_val_metrics_curves.png'))
    print('checkpoint', best_path)
   



if __name__ == "__main__":
    train_mintomics_classifier()