File size: 6,294 Bytes
5611f26 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | import argparse
import json
import logging
import os
import time
from dataset import MDataset
from network import FusionGraph
from sklearn.model_selection import KFold
import pandas as pd
import torch
import torch.nn as nn
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('--mode', type=str, default='11',
help="0 for off and 1 for on. First digit for seq, second for globf")
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=128,
help='embedding dimensions (default: 128)')
parser.add_argument('--dropout-ratio', type=float, dest='dropout_ratio', default=0.3,
help='dropout ratio (default: 0.3)')
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="gat",
help='gnn type (gin, gcn, gat, graphsage)')
parser.add_argument('--fusion', type=str, default='attention',
help='fusion type (attention, weighted, concat)')
# 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')
parser.add_argument('--loss', type=str, default='bce',
help='loss function')
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)
device = torch.device("cpu" if args.gpu == -1 or not torch.cuda.is_available() else f"cuda:{args.gpu}")
results = pd.DataFrame()
logging.info('Loading Test Dataset')
qlx_set = MDataset(threshold=args.threshold, mode='qlx', max_length=args.max_length, pdb_src=args.pdb_src)
qlx_loader = DataLoader(qlx_set, batch_size=args.batch_size, follow_batch=['x_s'], shuffle=False)
for fusion in ['attention', 'concat', 'weighted']:
weight_dir = "./run/" + '-'.join([args.mode, fusion, args.loss, str(args.batch_size), str(args.lr), str(args.epochs), args.pdb_src])
pred_all = []
gt_all = []
for i in range(1, 6):
model = FusionGraph(num_layer=args.num_layer, input_dim=43, emb_dim=args.emb_dim, out_dim=qlx_set.num_classes, JK="last",
drop_ratio=args.dropout_ratio, graph_pooling=args.graph_pooling, gnn_type=args.gnn_type, aux_mode=args.mode,
fusion_type=fusion)
model.load_state_dict(torch.load(os.path.join(weight_dir, f'model_{i}.pth')))
model.to(device).eval()
with torch.no_grad():
for data in qlx_loader:
data = data.to(device)
gt_all.append(torch.tensor(data.gt, device=device))
out = model(data)
pred_all.append(out)
pred_all = torch.nn.functional.sigmoid(torch.cat(pred_all, dim=0)).squeeze().cpu().numpy()
gt_all = torch.cat(gt_all, dim=0).int().squeeze().cpu().numpy()
results[fusion] = pred_all.ravel(order='F')
results['gt'] = gt_all.ravel(order='F')
results.to_csv("preds.csv", index=False)
if __name__ == "__main__":
main()
|