Test.py

import time
import argparse
import math
import numpy as np
import pandas as pd
import torch
from torch_geometric.loader import DataLoader
from tqdm import tqdm
from sklearn.metrics import accuracy_score, f1_score, matthews_corrcoef
from models.hyper.mhnn import MHNNM
from models.Mymodel import Mymodel
from mRNAdataset import mRNAdataset,mRNAdataset2,CombinedDataset
from utils import Logger, seed_everything
from utils.metrics import evaluate_metrics
import warnings
warnings.filterwarnings('ignore')
import logging
from utils.optimizers import *

def compute_acc(outputs, targets):
    preds = outputs
    count = 0
    k = 0
    for i in range(targets.shape[0]):
        p = sum(np.logical_and(targets[i], preds[i]))
        q = sum(np.logical_or(targets[i], preds[i]))
        if q == 0:
            k += 1
            continue
        count += p / q
    return count / (targets.shape[0] - k)


def evaluate_single_label_performance(model, test_loader, device, labels, output_file='performance_metrics.csv'):
    model.eval()
    all_preds = []
    all_labels = []

    with torch.no_grad():
        for batch in tqdm(test_loader, desc="Evaluating"):
            graph_data, seq_data = batch
            graph_data = graph_data.to(device)
            if isinstance(seq_data, dict):
                seq_data = {k: v.to(device) for k, v in seq_data.items()}
            else:
                seq_data = seq_data.to(device)
            outputs, _, _ = model(graph_data, seq_data, mask=None)
            preds = torch.sigmoid(outputs).cpu().numpy() > 0.5  # Binary classification threshold
            all_preds.append(preds)
            all_labels.append(graph_data.y.cpu().numpy())

    all_preds = np.concatenate(all_preds, axis=0)
    all_labels = np.concatenate(all_labels, axis=0)

    # Calculate metrics for each label
    results = []
    for idx, label in enumerate(labels):
        acc = accuracy_score(all_labels[:, idx], all_preds[:, idx])
        f1 = f1_score(all_labels[:, idx], all_preds[:, idx])
        mcc = matthews_corrcoef(all_labels[:, idx], all_preds[:, idx])
        results.append({
            'Label': label,
            'Accuracy': acc,
            'F1 Score': f1,
            'MCC': mcc
        })
        logging.info(f"Label: {label}")
        logging.info(f"Accuracy: {acc:.4f}")
        logging.info(f"F1 Score: {f1:.4f}")
        logging.info(f"MCC: {mcc:.4f}")

    # Save results to CSV file
    df = pd.DataFrame(results)
    df.to_csv(output_file, index=False)
    logging.info(f"Performance metrics saved to {output_file}")

    return results


if __name__ == '__main__':
    print('Task start time:')
    print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
    start_time = time.time()

    parser = argparse.ArgumentParser(description='Training')

    # Dataset arguments
    parser.add_argument('--data_dir', type=str, default='data/')

    # Training hyperparameters
    parser.add_argument('--seed', default=0, type=int)
    parser.add_argument('--device', type=int, default=0)
    parser.add_argument('--epochs', default=400, type=int)
    parser.add_argument('--batch_size', type=int, default=32)
    parser.add_argument('--lr', default=0.0001, type=float)
    parser.add_argument('--min_lr', default=0.000001, type=float)
    parser.add_argument('--wd', default=0.0, type=float)
    parser.add_argument('--clip_gnorm', default=None, type=float)
    parser.add_argument('--log_steps', type=int, default=1)
    parser.add_argument('--runs', default=1, type=int)

    # CNN hyperparameters
    parser.add_argument("--seq_max_len", default=6000, type=int)
    parser.add_argument("--seq_in_channels", default=5, type=int)
    parser.add_argument("--seq_gn", default=[0, 0.3, 0, 0.3, 0.3], nargs="+", type=float)
    parser.add_argument("--seq_kernel_size", default=15, type=int)
    parser.add_argument("--seq_conv_hidden", default=[128, 128, 256, 256, 256], nargs="+", type=int)
    parser.add_argument("--seq_maxpool", default=[2, 4, 2, 4, 4], nargs="+", type=int)
    parser.add_argument("--seq_lin_hidden", default=[128, 64], nargs="+", type=int)
    parser.add_argument("--seq_conv_dropout", default=0.2, type=float)
    parser.add_argument("--seq_lin_dropout", default=0.2, type=float)
    parser.add_argument("--seq_init", default="xavier_uniform", type=str)
    parser.add_argument('--pretrained_model', default='./DNAbert2_attention', type=str)
    parser.add_argument('--max_seq_len', default=6000, type=int, help='maximum sequence length for BERT input')
    parser.add_argument("--cl_n_layers", default=3, type=int)
    parser.add_argument("--cl_layers_size", default=[], nargs="+")
    parser.add_argument("--cl_dropout", default=0.1, type=float)
    parser.add_argument("--att_d", default=256, type=int)

    # Model hyperparameters
    parser.add_argument('--method', default='combine', help='model type')
    parser.add_argument('--All_num_layers', default=3, type=int, help='number of basic blocks')
    parser.add_argument('--MLP1_num_layers', default=2, type=int, help='layer number of mlps')
    parser.add_argument('--MLP2_num_layers', default=2, type=int, help='layer number of mlp2')
    parser.add_argument('--MLP3_num_layers', default=2, type=int, help='layer number of mlp3')
    parser.add_argument('--MLP4_num_layers', default=2, type=int, help='layer number of mlp4')
    parser.add_argument('--MLP_hidden', default=64, type=int, help='hidden dimension of mlps')
    parser.add_argument('--output_num_layers', default=2, type=int)
    parser.add_argument('--output_hidden', default=64, type=int)
    parser.add_argument('--aggregate', default='mean', choices=['sum', 'mean'])
    parser.add_argument('--normalization', default='ln', choices=['bn', 'ln', 'None'])
    parser.add_argument('--activation', default='relu', choices=['Id', 'relu', 'prelu'])
    parser.add_argument('--dropout', default=0.0, type=float)

    args = parser.parse_args()
    print(args)
    seed = args.seed
    seed_everything(seed=seed, workers=True)
    print(f'Seed: {seed}')
    device = f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu'
    device = torch.device(device)

    train_graph_dataset = mRNAdataset(root=args.data_dir, partition='train')
    train_seq_dataset = mRNAdataset2(root=args.data_dir, partition='train')
    valid_graph_dataset = mRNAdataset(root=args.data_dir, partition='val')
    valid_seq_dataset = mRNAdataset2(root=args.data_dir, partition='val')
    test_graph_dataset = mRNAdataset(root=args.data_dir, partition='test')
    test_seq_dataset = mRNAdataset2(root=args.data_dir, partition='test')

    # Combine datasets
    train_dataset = CombinedDataset(train_graph_dataset, train_seq_dataset)
    valid_dataset = CombinedDataset(valid_graph_dataset, valid_seq_dataset)
    test_dataset = CombinedDataset(test_graph_dataset, test_seq_dataset)

    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=False)
    valid_loader = DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=False)
    test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False)

    model = Mymodel(9, args)
    ema_model = ModelEma(model, decay=0.999)
    model_path = '/home/tu/hyper/model_best/best_ema_model.pth'
    state_dict = torch.load(model_path, map_location=torch.device('cpu'))
    ema_model.load_state_dict(state_dict)
    ema_model.to(args.device)

    # Optimize thresholds and evaluate performance
    logging.info("Optimizing thresholds and evaluating single-label performance...")
    unique_labels = ['Exosome', 'Nucleus', 'Nucleoplasm', 'Chromatin',
                     'Nucleolus', 'Cytosol', 'Membrane', 'Ribosome', 'Cytoplasm']
    results = evaluate_single_label_performance(ema_model, test_loader, device, unique_labels)

    # Log total execution time
    end_time = time.time()
    logging.info(f"Total execution time: {(end_time - start_time):.2f} seconds")