-
Notifications
You must be signed in to change notification settings - Fork 9
/
tuner.py
executable file
·221 lines (169 loc) · 8.89 KB
/
tuner.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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
# Pretrain a given model on NIH dataset
import argparse
import os
from utils import imgUtils, trainFeatures
from covid_models import hyperModel
import kerastuner
from kerastuner.tuners import BayesianOptimization
import pickle
"""Tune model
This script will use keras tuner to tune one of the six models used in ensemble
with a given dataset. The fully connected layers will be freezed at first. Then
all layers will be unfreezed. Hyperparameters can be changed in the main functions.
"""
def get_args():
"""
This function gets various user input form command line. The user input variables
include the name of the model to be tuned, the size of the input images,
the path to the iamge dataset, the path to the pretrianed weight file, and the
patht to the output directory where results will be saved.
Returns:
parser.parse_args() (list): a list of user input values.
"""
# Implement command line argument
parser = argparse.ArgumentParser(description='Use keras tuner to find best hyper parameter.')
parser.add_argument('-m', '--model', dest='model_name', metavar = 'model_name',
choices = ['ResNet-50', 'Xception', 'DenseNet-121', 'Inception-V3',
'Inception-ResNet-V2', 'EfficientNet-B2'],
type = str, nargs = 1, required = True,
help = 'the name of the model to be trained.\n Choose from ResNet-50, Xception, DenseNet-121, Inception-V3,'
'Inception-ResNet-V2, EfficientNet-B2')
parser.add_argument('--size', '-s', dest='img_size', metavar = 'img_size',
type = int, nargs = 1, required = True,
help = 'the size of dataset images')
parser.add_argument('--path', '-p', dest='path', metavar='DATA_path', type=str, nargs=1,
required = True, help='the path that contains the dataset.')
parser.add_argument('--weights', '-w', dest='weight_path', metavar='weight_path',
type=str, nargs=1, required=True, help='the path to pretrained weights, either NIH weight '
'file if training from scratch or corresponding weight '
'file from our pretrained weights if fine tuning '
'DeepCOVID-XR.')
parser.add_argument('--output', '-o', dest='output', metavar='prediction_output_path', type=str,
default=None, required=False, help='the directory to output best model weights and '
'hyperparameters; if not provided will output to current '
'working directory')
return parser.parse_args()
def make_path(data_dir, base, exp_name):
"""
This function creates path to save the training results and weights.
Parameters:
data_dir (string): the path to the parent directory of training and
validation datasets.
base (string): the path to the parent directory of saved weights.
exp_name (string): a unique name for different experiment runs.
Returns:
train_path (string): the path to the training dataset.
valid_path (string): the path to the validation dataset.
freeze_weight_save_path (string): the path to the trained weight with layers
freezed.
unfreeze_weight_save_path (string): the path to the trained weight with all
layers unfreezed.
freeze_img_save_path (string): the path to the result images with layers
freezed.
unfreeze_img_save_path (string): the path to the result images with all
layers unfreezed.
"""
train_path = os.path.join(data_dir, 'Train')
valid_path = os.path.join(data_dir, 'Validation')
if (not os.path.isdir(train_path)) or (not os.path.isdir(valid_path)):
print('Please split images into train directory and validation directory.')
exit()
if not os.path.isdir(os.path.join(base, 'tuner')):
os.mkdir(os.path.join(base, 'tuner'))
freeze_save_path = os.path.join(base, 'tuner', 'initial_cps', exp_name + '.h5')
unfreeze_save_path = os.path.join(base, 'tuner', 'cps', exp_name + '.h5')
best_model_path = os.path.join(base, 'tuner', 'best_models_and_params', exp_name, 'model')
best_weight_path = os.path.join(base, 'tuner', 'best_models_and_params', exp_name, 'model_weights.h5')
best_param_path = os.path.join(base, 'tuner', 'best_models_and_params', exp_name, 'model_params')
if not os.path.exists(os.path.dirname(freeze_save_path)):
os.makedirs(os.path.dirname(freeze_save_path), exist_ok=True)
if not os.path.exists(os.path.dirname(unfreeze_save_path)):
os.makedirs(os.path.dirname(unfreeze_save_path), exist_ok=True)
if not os.path.exists(best_model_path):
os.makedirs(best_model_path, exist_ok=True)
if not os.path.exists(os.path.dirname(best_weight_path)):
os.makedirs(os.path.dirname(best_weight_path), exist_ok=True)
if not os.path.exists(best_param_path):
os.makedirs(best_param_path, exist_ok=True)
return train_path, valid_path, freeze_save_path, unfreeze_save_path, best_model_path, best_weight_path, best_param_path
if __name__=='__main__':
"""
The main function sets vairous tuner parameters such as batch size and image
augmentation parameters. The fully connected layers are first trained for 50
epochs and then the entire network is trained for another 50 epochs.
The hyper parameters are set in this main function and can be changed below.
The best model hyperparameters will be pickled and saved. The best model and
weights will also be saved.
"""
batch_size = 16
rotation_range = 20
height_shift = 0.05
width_shift = 0.05
args = get_args()
data_path = os.path.normpath(args.path[0])
model_name = args.model_name[0]
img_size = args.img_size[0]
weights = os.path.normpath(args.weight_path[0])
output_path = args.output
if output_path is not None:
output_path = os.path.normpath(output_path)
else:
output_path = os.getcwd()
exp_name = model_name + '_' + str(img_size)
base_dir = output_path
train_dir, valid_dir, freeze_dir, unfreeze_dir, model_dir, weight_dir, param_dir = make_path(data_path, base_dir, exp_name)
img_proc = imgUtils(img_size)
train_idg, val_idg = img_proc.dataGen(rotation_range, height_shift, width_shift)
train_gen, val_gen = img_proc.generator(batch_size, train_idg, val_idg, train_dir, valid_dir)
lr = 0.001
momentum = 0.9
nestrov = True
patience_rlr = 3
patience_es = 5
factor = 0.1
min_delta = 0.001
monitor = 'val_auc'
epoch = 50
features = trainFeatures()
rlr = features.setRLP(monitor, factor, patience_rlr)
es = features.setES(monitor, patience_es, min_delta)
cp = features.setCP(monitor, freeze_dir)
freeze_model, model, base = features.getModel(model_name, img_size, weights)
features.compileModel(freeze_model, lr, momentum, nestrov)
model_history = features.generator(freeze_model, train_gen, val_gen, epoch, cp, rlr, es)
img_proc.plot_save(model_history, base_dir)
# Add dropout layer and run tuner with entire model
model = features.load(model, freeze_dir)
model = features.unfreeze(model)
patience_es = 10
es = features.setES(monitor, patience_es, min_delta)
cp = features.setCP(monitor, unfreeze_dir)
hp = hyperModel(base, freeze_dir)
TOTAL_TRIALS = 20
EXECUTION_PER_TRIAL = 1
EPOCHS = 50
tuner = BayesianOptimization(
hp,
max_trials=TOTAL_TRIALS,
objective=kerastuner.Objective("val_auc", direction="max"),
executions_per_trial=EXECUTION_PER_TRIAL,
directory=base_dir,
project_name=exp_name
)
history = tuner.search(train_gen,
epochs=EPOCHS,
validation_data=val_gen,
callbacks = [es, cp],
verbose =2,
use_multiprocessing=False)
# Save best model and weight
best_model = tuner.get_best_models()[0]
best_config = best_model.optimizer.get_config()
best_hyperparameters = tuner.get_best_hyperparameters()[0].get_config()
best_hyperparameters_values = tuner.get_best_hyperparameters()[0].values
best_model.save(model_dir)
best_model.save_weights(weight_dir)
with open(os.path.join(param_dir, 'hyperparameters.txt'), "w") as text_file:
text_file.write(str(best_hyperparameters))
pickle.dump(best_hyperparameters_values, open(os.path.join(param_dir,'hyperparameters.pickle'), 'wb'))
print('Done')