import argparse import json import logging import os import time from dataset import SDataset, MDataset from network import MMGraph from sklearn.model_selection import KFold from train import train import torch import torch.nn as nn from torch.utils.data import Subset from torch_geometric.loader import 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)") parser.add_argument('--num-layer', type=int, dest='num_layer', default=2, help='number of GNN message passing layers (default: 2)') parser.add_argument('--emb-dim', type=int, dest='emb_dim', default=90, help='embedding dimensions (default: 200)') parser.add_argument('--dropout-ratio', type=float, dest='dropout_ratio', default=0.5, help='dropout ratio (default: 0.5)') parser.add_argument('--graph-pooling', type=str, dest='graph_pooling', default="attention", help='graph level pooling (sum, mean, max, attention)') parser.add_argument('--gnn-type', type=str, dest='gnn_type', default="gatv2", help='gnn type (gin, gcn, gat, graphsage)') # 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=110, help='number of epochs to train (default: 50)') parser.add_argument('--lr', type=float, default=0.002, 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, follow_batch=['x_s'], shuffle=True, drop_last=True) valid_loader = DataLoader(valid_set, batch_size=args.batch_size, follow_batch=['x_s'], shuffle=False) qlx_loader = DataLoader(qlx_set, batch_size=1, follow_batch=['x_s'], shuffle=False) saap_loader = DataLoader(saap_set, batch_size=1, follow_batch=['x_s'], shuffle=False) if args.model == 'mm': model = MMGraph(num_layer=args.num_layer, input_dim=43, emb_dim=args.emb_dim, out_dim=set_all.num_classes, JK="last", drop_ratio=args.dropout_ratio, graph_pooling=args.graph_pooling, gnn_type=args.gnn_type, max_length=args.max_length) 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()