-
Notifications
You must be signed in to change notification settings - Fork 0
/
main_model_selection.py
108 lines (84 loc) · 4.91 KB
/
main_model_selection.py
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
import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""
import torch
from sklearn.metrics import precision_score, recall_score, confusion_matrix, f1_score, accuracy_score
from data import FingerprintDataset, CompoundDataset
from train.train_assess import cross_validation
from train.exp_NeFPNN import run_NeFPNN
from train.exp_MoRF import run_MoRF
import train.utils as utils
from models import *
import argparse
import warnings
import time
random_seed = 42
torch.manual_seed(random_seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(random_seed)
def run(args):
print("Loading DataSets...")
start = time.time()
adaptive = (args.model_name != 'MoNN' and args.model_name != 'MoRF')
# Load train and test datasets
if adaptive:
# Dataset for Adaptive fingerprint
# It contains the atom/bond features used in https://github.com/XuhanLiu/NGFP + dose
train_dataset = CompoundDataset(['./Datasets/NGFP_train_dataset.csv', './Datasets/NGFP_train_atom.csv', './Datasets/NGFP_train_bond.csv'], bidir=True)
test_dataset = CompoundDataset(['./Datasets/NGFP_test_dataset.csv', './Datasets/NGFP_test_atom.csv', './Datasets/NGFP_test_bond.csv'], bidir=True)
else:
# Dataset for static fingerprint
# It contains precomputed Extended-Connectivity Fingerprint + dose
train_dataset = FingerprintDataset(args.MoNN_tr_path)
test_dataset = FingerprintDataset(args.MoNN_tr_path.replace('train', 'test'))
t = time.time() - start
print('\t %d hours, %d minutes, %d seconds, %.3d ms' % (t // 3600, t % 3600 // 60, t % 60, (t % 1) * 1000))
print("Building model <model>...")
start = time.time()
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
# Define optimizers
optims = utils.get_optims_confs()
if 'MoNN' in args.model_name:
inp_dim = len(train_dataset[0].x[0])
models = utils.get_MoNN_confs(inp_dim)
else:
inp_dim = train_dataset[0].x.size()[1]
dose_ft_dim = train_dataset[0].graph_features.size()[1]
edge_ft_dim = train_dataset[0].edge_attr.size()[1]
if 'LinNN' in args.model_name:
models = utils.get_LinNN_confs(inp_dim, dose_ft_dim, device=device)
else:
models = utils.get_DGN_confs(args.model_name, inp_dim, dose_ft_dim, edge_ft_dim, device=device)
loss_fun = torch.nn.CrossEntropyLoss()
t = time.time() - start
print('\t %d hours, %d minutes, %d seconds, %.3d ms' % (t // 3600, t % 3600 // 60, t % 60, (t % 1) * 1000))
metrics = [utils.multiclass_roc_auc_score, utils.per_class_accuracy, f1_score, recall_score,
precision_score, accuracy_score, confusion_matrix]
if args.num_process > 1 and torch.cuda.is_available():
args.num_process = 1
warnings.warn("num_process was decreased to 1 because possible errors with cuda", RuntimeWarning)
cross_validation(args.k, models, optims, train_dataset, loss_fun, batch_size=args.batch,
early_stop=True, epochs=args.e, name=args.exp_name, metrics=metrics, complex_cv=args.simple_cv,
n_workers=args.num_process, random_seed=random_seed, draw_plot=True, device=device)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--model_name', help='Name of the model to run', type=str, default='GaNN', choices=utils.conf.keys())
parser.add_argument('--e', help='The number of epochs', type=int, default=900)
parser.add_argument('--batch', help='The batch size', type=int, default=512)
parser.add_argument('--k', help='The number of inner folds in the k-fold/nested cross-validation', type=int, default=3)
parser.add_argument('--simple_cv', help='The type of cross-validation. If True the simple cross-validation is performed, else complex cross-validatio is perfomed.', action='store_true')
parser.add_argument('--MoNN_tr_path', help='The path of the MoNN training dataset', type=str, default='') # example: ./Datasets/static_fp_complex_cv/train_random_forest_group-FTT-L1024-R3.csv
parser.add_argument('--exp_name', help='The name of the experiment', type=str, default="3-cv-GaNN")
parser.add_argument('--num_process', help='The number of concurrent jobs', type=int, default=15)
args = parser.parse_args()
beginning = time.time()
if args.model_name == 'MoRF':
run_MoRF('cv', random_seed, args)
elif args.model_name == 'NeFPNN':
optims = utils.get_optims_confs()
model_confs = utils.get_DGN_confs('NeFPNN', 128, 1)
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
run_NeFPNN('cv', optims, model_confs, random_seed, device, args)
else:
run(args)
t = time.time() - beginning
print('Required time for main.py in ' + args.exp_name + ' mode: %d hours, %d minutes, %d seconds, %.3d ms' % (t // 3600, t % 3600 // 60, t % 60, (t % 1) * 1000))