import argparse import json import logging import os import time from dataset import SDataset, MDataset from network import MMPeptide, SEQPeptide, VoxPeptide, MMFPeptide, SMPeptide 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='mm', help='model resnet26, bi-gru') 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('--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=64, help="MIC threshold for determine labels (default: 64)") # 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=128, 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.02, help='learning rate (default: 0.002)') 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='micro', help='metric average type') # args for losses parser.add_argument('--loss', type=str, default='mlce', 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 weight_dir = "./run/" + args.task_type + '-' + args.task + "-m-" + args.model + '-' + args.loss + str(args.batch_size) + str(args.lr) + str(args.epochs) + args.pdb_src if not os.path.exists(weight_dir): raise ValueError logging.basicConfig(handlers=[ # logging.FileHandler(filename=os.path.join(weight_dir, "eval.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) 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) # 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) if args.model == 'seq': model = SEQPeptide(classes=set_all.num_classes, q_encoder='mlp') elif args.model == 'voxel': model = VoxPeptide(classes=set_all.num_classes) elif args.model == 'mm': # model = MMPeptide(classes=train_set.num_classes, q_encoder='tf', ) # attention='hamburger' # model = SMPeptide(classes=train_set.num_classes, q_encoder='mlp', max_length=30) # attention='hamburger' model = MMPeptide(classes=set_all.num_classes, q_encoder='mlp', max_length=args.max_length)# , attention='hamburger' elif args.model == 'mmf': model = MMFPeptide(classes=set_all.num_classes, q_encoder='mlp', max_length=args.max_length) # attention='hamburger' 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) weights_path = f"{weight_dir}/model_{fold + 1}.pth" model.load_state_dict(torch.load(weights_path)) # early_stopping = EarlyStopping(patience=args.patience, path=weights_path) logging.info(f'Running Cross Validation {fold + 1}') logging.info(f'Fold {fold + 1} Valid set:{len(valid_set)}, Test set:qlx {len(qlx_set)} saap {len(saap_set)}') if args.task_type in ('mlc', 'slc') : _, ap, f1, acc, auc = train(args, None, model, None, valid_loader, device, None, None) logging.info(f'ap: {ap:.3f}, f1: {f1:.3f}, acc: {acc:.3f}, auc: {auc:.3f}') best_perform_list[fold] = np.asarray([ap, f1, acc, auc]) _, qlx_ap, qlx_f1, qlx_acc, qlx_auc = train(args, None, model, None, qlx_loader, device, None, None) logging.info(f'QLX results, ap: {qlx_ap:.3f}, f1: {qlx_f1:.3f}, acc: {qlx_acc:.3f}, auc: {qlx_auc:.3f}') qlx_perform_list[fold] = np.asarray([qlx_ap, qlx_f1, qlx_acc, qlx_auc]) _, saap_ap, saap_f1, saap_acc, saap_auc = train(args, None, model, None, saap_loader, device, None, None) logging.info(f'SAAP results, ap: {saap_ap:.3f}, f1: {saap_f1:.3f}, acc: {saap_acc:.3f}, auc: {saap_auc:.3f}') saap_perform_list[fold] = np.asarray([saap_ap, saap_f1, saap_acc, saap_auc]) else: raise NotImplementedError 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+'/eval.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()