t3.1/main.py

import argparse
import json
import logging
import os
import time

from dataset import MDataset
from network import SEQPeptide
from sklearn.model_selection import KFold
from train import train
import torch
import torch.nn as nn
from torch.utils.data import Subset, DataLoader
import numpy as np
from loss import MLCE, SuperLoss, LogCoshLoss
from utils import set_seed


parser = argparse.ArgumentParser(description='resnet26')
# model setting
parser.add_argument('--model', type=str, default='mlp',
                    help='mlp gru lstm rnn tf mamba')
parser.add_argument('--fusion', type=str, default='1',
                    help="Seed for splitting dataset (default 1)")

# task & dataset setting
parser.add_argument('--pdb-src', type=str, dest='pdb_src', default='af',
                    help='af or hf')
parser.add_argument('--task-type', type=str, dest='task_type', default='mlc',
                    help='mlc or slc')
parser.add_argument('--data-ver', type=str, dest='data_ver', default='0920',
                    help='data version')
parser.add_argument('--task', type=str, default='all',
                    help='task: anti toxin anti-all mechanism anti-binary anti-regression mic')
parser.add_argument('--classes', type=int, default=6,
                    help='model')
parser.add_argument('--max-length', dest='max_length', type=int, default=30,
                    help='Max length for sequence filtering')
parser.add_argument('--split', type=int, default=5,
                    help="Split k fold in cross validation (default: 5)")
parser.add_argument('--seed', type=int, default=1,
                    help="Seed for splitting dataset (default: 1)")
parser.add_argument('--threshold', type=float, default=128,
                    help="MIC threshold for determine labels (default: 128)")

# training setting
parser.add_argument('--gpu', type=int, default=0,
                    help='GPU index to use, -1 for CPU (default: 0)')
parser.add_argument('--batch-size', type=int, dest='batch_size', default=256,
                    help='input batch size for training (default: 128)')
parser.add_argument('--epochs', type=int, default=50,
                    help='number of epochs to train (default: 50)')
parser.add_argument('--lr', type=float, default=0.001,
                    help='learning rate (default: 0.001)')
parser.add_argument('--decay', type=float, default=0.0005,
                    help='weight decay (default: 0.0005)')
parser.add_argument('--warm-steps', type=int, dest='warm_steps', default=0,
                    help='number of warm start steps for learning rate (default: 10)')
parser.add_argument('--patience', type=int, default=10,
                    help='patience for early stopping (default: 10)')
parser.add_argument('--pretrain', type=str, dest='pretrain', default='',
                    help='path of the pretrain model')  # /home/duadua/Desktop/fetal/3dpretrain/runs/e50.pth
parser.add_argument('--metric-avg', type=str, dest='metric_avg', default='macro',
                    help='metric average type')

# args for losses
parser.add_argument('--loss', type=str, default='bce',
                    help='loss function (mlce, sl, mix)')

parser.add_argument('--bias-curri', dest='bias_curri', action='store_true', default=False,
                    help='directly use loss as the training data (biased) or not (unbiased)')
parser.add_argument('--anti-curri', dest='anti_curri', action='store_true', default=False,
                    help='easy to hard (curri), hard to easy (anti)')
parser.add_argument('--std-coff', dest='std_coff', type=float, default=1,
                    help='the hyper-parameter of std')

args = parser.parse_args()


def main():
    set_seed(args.seed)

    if args.task_type == 'slc':
        if args.task == 'all':
            raise ValueError('Choose one task number to run single label classification')
        args.classes = 1
    elif args.task_type == 'mlc':
        pass
    else:
        raise NotImplementedError

    if args.loss == 'mlce' and args.task_type != 'slc':
        criterion = MLCE()
    elif args.loss == "bce" or args.task_type == 'slc':
        args.loss = "bce"
        criterion = nn.BCEWithLogitsLoss()
    else:
        criterion = 0
        pass
    weight_dir = "./run/" + args.model + '-' + args.loss + str(args.batch_size) + '-' + str(args.lr) + '-' + str(args.epochs)
    if not os.path.exists(weight_dir):
        os.makedirs(weight_dir)
    
    logging.basicConfig(handlers=[
        logging.FileHandler(filename=os.path.join(weight_dir, "training.log"), encoding='utf-8', mode='w+'),
        logging.StreamHandler()],
        format="%(asctime)s: %(message)s", datefmt="%F %T", level=logging.INFO)

    logging.info(f'saving_dir: {weight_dir}')
    
    with open(os.path.join(weight_dir, "config.json"), "w") as f:
        f.write(json.dumps(vars(args)))

    device = torch.device("cpu" if args.gpu == -1 or not torch.cuda.is_available() else f"cuda:{args.gpu}")

    logging.info('Loading Training Dataset')
    if args.task_type == 'mlc':
        set_all = MDataset(threshold=args.threshold, mode='train', max_length=args.max_length, pdb_src=args.pdb_src, data_ver=args.data_ver)
    # else:
    #     set_all = SDataset(threshold=args.threshold, mode='train', task=args.task, max_length=args.max_length, pdb_src=args.pdb_src)
    logging.info('Loading Test Dataset')
    if args.task_type == 'mlc':
        qlx_set = MDataset(threshold=args.threshold, mode='qlx', max_length=args.max_length, pdb_src=args.pdb_src)
        saap_set = MDataset(threshold=args.threshold, mode='saap', max_length=args.max_length, pdb_src=args.pdb_src)
    # else:
    #     qlx_set = SDataset(threshold=args.threshold, mode='qlx', task=args.task, max_length=args.max_length, pdb_src=args.pdb_src)
    #     saap_set = SDataset(threshold=args.threshold, mode='saap', task=args.task, max_length=args.max_length, pdb_src=args.pdb_src)

    best_perform_list = [[] for i in range(5)]
    qlx_perform_list = [[] for i in range(5)]
    saap_perform_list = [[] for i in range(5)]

    kf = KFold(n_splits=5, shuffle=True, random_state=42)

    for fold, (train_idx, val_idx) in enumerate(kf.split(set_all)):
        train_set= Subset(set_all, train_idx)
        valid_set = Subset(set_all, val_idx)
        # if args.loss == 'bce':
        #     weight = torch.Tensor([0] * args.classes)
        #     for i in train_set:
        #         weight += i[2]
        #     weight = (len(train_set) - weight) / weight
        #     criterion.register_buffer('pos_weight', weight.to(device))

        train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, drop_last=True)
        valid_loader = DataLoader(valid_set, batch_size=args.batch_size, shuffle=False)
        qlx_loader = DataLoader(qlx_set, batch_size=1, shuffle=False)
        saap_loader = DataLoader(saap_set, batch_size=1, shuffle=False)
        model = SEQPeptide(classes=set_all.num_classes, q_encoder=args.model)
        if len(args.pretrain) != 0:
            logging.info('loading pretrain model')
            # model = load_pretrain_model(model, torch.load(args.pretrain))
            model_state = model.state_dict()
            pretrained_state = torch.load(args.pretrain)
            pretrained_state = {k: v for k, v in pretrained_state.items() if
                                k in model_state and v.size() == model_state[k].size()}
            model_state.update(pretrained_state)
            model.load_state_dict(model_state)
            # model.load_state_dict(torch.load(args.pretrain), strict=False)
        model.to(device)
        # optimizer = torch.optim.Adam(model.parameters())
        optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
        # optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=True, weight_decay=5e-5)
        weights_path = f"{weight_dir}/model_{fold + 1}.pth"
        # early_stopping = EarlyStopping(patience=args.patience, path=weights_path)
        logging.info(f'Running Cross Validation {fold + 1}')
        logging.info(f'Fold {fold + 1}  Train set:{len(train_set)}, Valid set:{len(valid_set)}, Test set:qlx {len(qlx_set)} saap {len(saap_set)}')
        best_metric = 0
        best_qlx = 0
        best_saap = 0
        start_time = time.time()

        for epoch in range(1, args.epochs + 1):
            if args.task_type in ('mlc', 'slc') :
                train_loss, ap, f1, acc, auc = train(args, epoch, model, train_loader, valid_loader, device, criterion, optimizer)
                logging.info(f'Epoch: {epoch:03d}, Train Loss: {train_loss:.3f}, ap: {ap:.3f}, f1: {f1:.3f}, acc: {acc:.3f}, auc: {auc:.3f}')
                avg_metric = ap + f1 + acc + auc
                if avg_metric > best_metric:
                    logging.info(f'Epoch: {epoch:03d} New Best validation metrics, running test session')
                    torch.save(model.state_dict(), weights_path)
                    best_metric = avg_metric
                    best_perform_list[fold] = np.asarray([ap, f1, acc, auc])
                    
                    _, qlx_ap, qlx_f1, qlx_acc, qlx_auc = train(args, epoch, model, None, qlx_loader, device, None, None)
                    logging.info(f'Epoch: {epoch:03d} QLX results, ap: {qlx_ap:.3f}, f1: {qlx_f1:.3f}, acc: {qlx_acc:.3f}, auc: {qlx_auc:.3f}')
                    qlx_metric = qlx_ap + qlx_f1 + qlx_acc + qlx_auc
                    if qlx_metric > best_qlx:
                        best_qlx = qlx_metric
                        qlx_perform_list[fold] = np.asarray([qlx_ap, qlx_f1, qlx_acc, qlx_auc])
                    
                    _, saap_ap, saap_f1, saap_acc, saap_auc = train(args, epoch, model, None, saap_loader, device, None, None)
                    logging.info(f'Epoch: {epoch:03d} SAAP results, ap: {saap_ap:.3f}, f1: {saap_f1:.3f}, acc: {saap_acc:.3f}, auc: {saap_auc:.3f}')
                    saap_metric = saap_ap + saap_f1 + saap_acc + saap_auc
                    if saap_metric > best_saap:
                        best_saap = saap_metric
                        saap_perform_list[fold] = np.asarray([saap_ap, saap_f1, saap_acc, saap_auc])
                           
            else:
                raise NotImplementedError
            
        logging.info(f'used time {(time.time()-start_time)/3600:.2f}h')

    logging.info(f'Cross Validation Finished!')
    best_perform_list = np.asarray(best_perform_list)
    qlx_perform_list = np.asarray(qlx_perform_list)
    saap_perform_list = np.asarray(saap_perform_list)
    logging.info('Best validation perform list\n%s', best_perform_list)
    logging.info('mean: %s', np.round(np.mean(best_perform_list, 0), 3))
    logging.info('std: %s', np.round(np.std(best_perform_list, 0), 3))
    logging.info('Best qlx perform list\n%s', qlx_perform_list)
    logging.info('mean: %s', np.round(np.mean(qlx_perform_list, 0), 3))
    logging.info('std: %s', np.round(np.std(qlx_perform_list, 0), 3))
    logging.info('Best saap perform list\n%s', saap_perform_list)
    logging.info('mean: %s', np.round(np.mean(saap_perform_list, 0), 3))
    logging.info('std: %s', np.round(np.std(saap_perform_list, 0), 3))
    perform = open(weight_dir+'/result.txt', 'w')
    perform.write('Valid\n')
    perform.write(','.join([str(i) for i in np.mean(best_perform_list, 0)])+'\n')
    perform.write(','.join([str(i) for i in np.std(best_perform_list, 0)])+'\n')
    perform.write('qlx\n')
    perform.write(','.join([str(i) for i in np.mean(qlx_perform_list, 0)])+'\n')
    perform.write(','.join([str(i) for i in np.std(qlx_perform_list, 0)])+'\n')
    perform.write('saap\n')
    perform.write(','.join([str(i) for i in np.mean(saap_perform_list, 0)])+'\n')
    perform.write(','.join([str(i) for i in np.std(saap_perform_list, 0)])+'\n')


if __name__ == "__main__":
    main()