-
Notifications
You must be signed in to change notification settings - Fork 2
/
main.py
77 lines (67 loc) · 1.8 KB
/
main.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
import argparse
import json
from core.model import MedicalImageClassifier
import os
import tensorflow as tf
import faulthandler
def get_parser():
# create parser object
parser = argparse.ArgumentParser(
description='Neural network classifier for medical imaging using Tensorflow',
epilog='For questions and bug reports, contact Jacky Ko <jkmailbox1991@gmail.com>')
# add arguments
parser.add_argument(
'-v', '--verbose',
dest='verbose',
help='Show verbose output',
action='store_true')
parser.add_argument(
'-p','--phase',
dest='phase',
help='Training phase (default=TRAIN)',
choices=['TRAIN','PREDICT'],
default='TRAIN',
metavar='[TRAIN PREDICT]')
parser.add_argument(
'--config_json',
dest='config_json',
help='JSON file for model configuration',
type=str,
default='config.json',
metavar='FILENAME'
)
parser.add_argument(
'--gpu',
dest='gpu',
default='0',
type=str,
help='Select GPU device(s) (default = 0)',
metavar='GPU_IDs')
args = parser.parse_args()
if args.verbose:
args_dict = vars(args)
for key in sorted(args_dict):
print("{} = {}".format(str(key), str(args_dict[key])))
return args
def main(args):
# select gpu
# os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu) # e.g. "0,1,2", "0,2"
# read configuration file
with open(args.config_json) as config_json:
config = json.load(config_json)
# session config
config_proto = tf.ConfigProto()
config_proto.gpu_options.allow_growth = True
# config_proto.log_device_placement = True
with tf.Session(config=config_proto) as sess:
model = MedicalImageClassifier(sess,config)
if args.phase == "TRAIN":
model.train()
else:
model.predict()
if __name__=="__main__":
# for seg fault trace
faulthandler.enable()
args = get_parser()
main(args)