forked from chemprop/chemprop
-
Notifications
You must be signed in to change notification settings - Fork 0
/
sklearn_train.py
37 lines (31 loc) · 1.64 KB
/
sklearn_train.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
from argparse import ArgumentParser
from chemprop.parsing import add_train_args, modify_train_args
from chemprop.sklearn_train import cross_validate_sklearn
from chemprop.utils import create_logger
if __name__ == '__main__':
parser = ArgumentParser()
add_train_args(parser)
parser.add_argument('--class_weight', type=str,
choices=['balanced'],
help='How to weight classes (None means no class balance)')
parser.add_argument('--single_task', action='store_true', default=False,
help='Whether to run each task separately (needed when dataset has null entries)')
parser.add_argument('--radius', type=int, default=2,
help='Morgan fingerprint radius')
parser.add_argument('--num_bits', type=int, default=2048,
help='Number of bits in morgan fingerprint')
parser.add_argument('--model_type', type=str, choices=['random_forest', 'svm'], required=True,
help='scikit-learn model to use')
parser.add_argument('--num_trees', type=int, default=500,
help='Number of random forest trees')
args = parser.parse_args()
modify_train_args(args)
logger = create_logger(name='sklearn-train', save_dir=args.save_dir, quiet=args.quiet)
if args.metric is None:
if args.dataset_type == 'regression':
args.metric = 'rmse'
elif args.dataset_type == 'classification':
args.metric = 'auc'
else:
raise ValueError(f'Default metric not supported for dataset_type "{args.dataset_type}"')
cross_validate_sklearn(args, logger)