import argparse import json import logging import os import time from dataset import PeptidePairDataset, PeptidePairPicDataset, SimplePairClsDataset from network import DMutaPeptide, DMutaPeptideCNN#, DMutaPeptideWiden from sklearn.model_selection import KFold from train import train_cls import torch import torch.nn as nn from torch.utils.data import DataLoader, WeightedRandomSampler, RandomSampler, Subset import numpy as np from loss import MLCE, SuperLoss, LogCoshLoss, BMCLoss from utils import set_seed parser = argparse.ArgumentParser(description='resnet26') # model setting parser.add_argument('--model', type=str, default='resnet34', help='resnet34 resnet50 densenet') parser.add_argument('--q-encoder', dest='q_encoder', type=str, default='cnn', help='lstm mamba mla') parser.add_argument("--side-enc", dest='side_enc', type=str, default='lstm', help="use side features") parser.add_argument('--channels', type=int, default=16) parser.add_argument('--fusion', type=str, default='att', help='mlp att diff') parser.add_argument('--glob-feat', dest='glob_feat', action='store_true', default=False, help="use global features") parser.add_argument('--non-siamese', dest='non_siamese', action='store_true', default=False, help="use non-siamese architecture") parser.add_argument('--widen', action='store_true', default=False, help='use widen non-siamese architecture') # task & dataset setting parser.add_argument('--task', type=str, default='cls', help='reg or cls') parser.add_argument('--pdb-src', type=str, dest='pdb_src', default='af', help='af or hf') parser.add_argument('--data-ver', type=str, dest='data_ver', default='250228', help='data version') parser.add_argument('--one-way', action='store_true', dest='one_way', default=True, help='use one-way constructed dataset') 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('--run-folds', type=int, dest='run_folds', nargs='+', default=-1, help='specify which folds to run') parser.add_argument('--seed', type=int, default=1, help="Seed (default: 1)") parser.add_argument('--pcs', action='store_true', default=False, help='Consider protease cut site') parser.add_argument('--mix-pcs', dest='mix_pcs', action='store_true', default=False, help='Consider protease cut site') parser.add_argument('--resize', type=int, default=[768], nargs='+', help='resize the image') parser.add_argument('--llm-data', action='store_true', default=False, help='Use LLM augmentation data') # 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=32, help='input batch size for training (default: 128)') parser.add_argument('--epochs', type=int, default=50, help='number of epochs to train (default: 100)') 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') parser.add_argument('--metric-avg', type=str, dest='metric_avg', default='macro', help='metric average type') parser.add_argument('--loss', type=str, default='ce', help='loss function') parser.add_argument('--dir', action='store_true', default=False, help='use DIR') 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') parser.add_argument('--ft-epochs', dest='ft_epochs', type=int, default=15, help='fine-tune epochs') parser.add_argument('--ft-lr', dest='ft_lr', type=float, default=0.0002, help='fine-tune learning rate') parser.add_argument('--simple', dest='simple', action='store_true', default=False) args = parser.parse_args() if args.llm_data: args.simple = True if args.simple: args.one_way = True if args.run_folds == -1: args.run_folds = list(range(args.split)) def main(): set_seed(args.seed) if args.task == 'reg': args.classes = 1 if args.loss == "mse" or args.loss in ['ce']: args.loss = 'mse' criterion = nn.MSELoss() elif args.loss == "smoothl1": criterion = nn.SmoothL1Loss() elif args.loss == "super": criterion = SuperLoss() elif args.loss in ["bmc", "bmc_ln"]: criterion = BMCLoss() else: raise NotImplementedError("unimplemented regression task loss function") elif args.task == 'cls': args.classes = 2 if args.loss == 'ce' or args.loss in ['mse', 'smoothl1', 'super']: args.loss = 'ce' criterion = nn.CrossEntropyLoss() else: raise NotImplementedError("unimplemented classification task loss function") else: raise NotImplementedError("unimplemented task") if args.q_encoder in ['cnn', 'rn18']: weight_dir = f'./run-{args.task}/{args.q_encoder}{f"-non-siamese" if args.non_siamese else ""}-{args.fusion}-{args.channels}{f"-{args.side_enc}" if args.side_enc else ""}{"-mixpcs" if args.mix_pcs else ""}{"-pcs" if args.pcs==True else ""}{"-simple" if args.simple else ""}{"-llm" if args.llm_data else ""}{"-" + "x".join(str(n) for n in args.resize) if args.resize else ""}{"-gf" if args.glob_feat else ""}{"-oneway" if args.one_way else ""}-{args.loss + "-dir" if args.dir else args.loss}-{str(args.batch_size)}-{str(args.lr)}-{str(args.epochs)}' else: weight_dir = f'./run-{args.task}/{args.q_encoder}{f"-non-siamese" if args.non_siamese else ""}-{args.fusion}-{args.channels}{"-simple" if args.simple else ""}{"-llm" if args.llm_data else ""}{"-gf" if args.glob_feat else ""}{"-oneway" if args.one_way else ""}-{args.loss + "-dir" if args.dir else args.loss}-{str(args.batch_size)}-{str(args.lr)}-{str(args.epochs)}' logging.basicConfig(handlers=[ logging.FileHandler(filename=os.path.join(weight_dir, "finetune.log"), encoding='utf-8', mode='w+'), logging.StreamHandler()], format="%(asctime)s: %(message)s", datefmt="%F %T", level=logging.INFO) logging.info(f'Finetuning: {weight_dir}') device = torch.device("cpu" if args.gpu == -1 or not torch.cuda.is_available() else f"cuda:{args.gpu}") logging.info(f'Loading Training Dataset') train_set = SimplePairClsDataset(pad_length=args.max_length, ftr2=True, gf=args.glob_feat, q_encoder=args.q_encoder, side_enc=args.side_enc, pcs=args.pcs, resize=args.resize) logging.info('Loading Test Dataset') if args.q_encoder in ['cnn', 'rn18']: test_set = PeptidePairPicDataset(mode='r2_case', pad_length=args.max_length, task=args.task, gf=args.glob_feat, side_enc=args.side_enc, pcs=args.pcs, resize=args.resize) else: test_set = PeptidePairDataset(mode='r2_case', pad_length=args.max_length, task=args.task, gf=args.glob_feat) train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=8, pin_memory=True) test_loader = DataLoader(test_set, batch_size=args.batch_size, shuffle=False, num_workers=8, pin_memory=True) best_perform_list = [[] for i in range(5)] for fold in range(args.split): logging.info(f'Finetuning Fold {fold}') logging.info(f'Fold {fold} Train set:{len(train_set)}, Test set: {len(test_set)}') # if args.widen: # model = DMutaPeptideWiden(q_encoder=args.q_encoder, classes=args.classes, channels=args.channels, dir=args.dir, gf=args.glob_feat, fusion=args.fusion, side_enc=args.side_enc) # else: if args.q_encoder in ['cnn', 'rn18']: model = DMutaPeptideCNN(q_encoder=args.q_encoder, classes=args.classes, channels=args.channels, dir=args.dir, gf=args.glob_feat, side_enc=args.side_enc, fusion=args.fusion, non_siamese=args.non_siamese) else: model = DMutaPeptide(q_encoder=args.q_encoder, classes=args.classes, channels=args.channels, dir=args.dir, gf=args.glob_feat, fusion=args.fusion, non_siamese=args.non_siamese) weights_path = f"{weight_dir}/model_{fold}.pth" model.to(device) # model.load_state_dict(torch.load(weights_path.replace('.pth', '_test.pth'), map_location=device), strict=False) model.load_state_dict(torch.load(weights_path, map_location=device), strict=False) optimizer = torch.optim.AdamW(model.parameters(), lr=args.ft_lr) best_metric = -float('inf') if args.task == 'cls': for epoch in range(1, args.ft_epochs + 1): train_loss, ap, auc, f1, acc = train_cls(args, epoch, model, train_loader, test_loader, device, criterion, optimizer) logging.info(f'Epoch: {epoch:03d} Train Loss: {train_loss:.3f}, ap: {ap:.3f}, auc: {auc:.3f}, f1: {f1:.3f}, acc: {acc:.3f}') avg_metric = ap + auc #+ f1 + acc if avg_metric > best_metric: logging.info(f'Epoch: {epoch:03d} New best VALIDATION metrics') best_metric = avg_metric best_perform_list[fold] = np.asarray([ap, auc, f1, acc]) torch.save(model.state_dict(), weights_path.replace('.pth', '_ft.pth')) if __name__ == "__main__": main()