-
Notifications
You must be signed in to change notification settings - Fork 74
/
keras_train.py
executable file
·153 lines (136 loc) · 7.87 KB
/
keras_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
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
import tensorflow.python as tf
from tensorflow.contrib.data import assert_element_shape
from tensorflow.python import keras
from tensorflow.python.keras.callbacks import LearningRateScheduler
from tensorflow.python.keras.callbacks_v1 import TensorBoard
from tools.utils import Helper, create_loss_fn, INFO, ERROR, NOTE
from tools.custom import Yolo_Precision, Yolo_Recall
from models.yolonet import *
import os
from pathlib import Path
from datetime import datetime
import numpy as np
import sys
import argparse
from termcolor import colored
from tensorflow_model_optimization.python.core.api.sparsity import keras as sparsity
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
keras.backend.set_session(sess)
def write_arguments_to_file(args, filename):
with open(filename, 'w') as f:
for key, value in vars(args).items():
f.write('%s: %s\n' % (key, str(value)))
def main(args, train_set, class_num, pre_ckpt, model_def,
depth_multiplier, is_augmenter, image_size, output_size,
batch_size, rand_seed, max_nrof_epochs, init_learning_rate,
learning_rate_decay_factor, obj_weight, noobj_weight,
wh_weight, obj_thresh, iou_thresh, vaildation_split, log_dir,
is_prune, initial_sparsity, final_sparsity, end_epoch, frequency):
# Build path
log_dir = (Path(log_dir) / datetime.strftime(datetime.now(), '%Y%m%d-%H%M%S')) # type: Path
ckpt_weights = log_dir / 'yolo_weights.h5'
ckpt = log_dir / 'yolo_model.h5'
if not log_dir.exists():
log_dir.mkdir(parents=True)
write_arguments_to_file(args, str(log_dir / 'args.txt'))
# Build utils
h = Helper(f'data/{train_set}_img_ann.npy', class_num, f'data/{train_set}_anchor.npy',
np.reshape(np.array(image_size), (-1, 2)), np.reshape(np.array(output_size), (-1, 2)), vaildation_split)
h.set_dataset(batch_size, rand_seed, is_training=(is_augmenter == 'True'))
# Build network
network = eval(model_def) # type :yolo_mobilev2
yolo_model, yolo_model_warpper = network([image_size[0], image_size[1], 3], len(h.anchors[0]), class_num, alpha=depth_multiplier)
if pre_ckpt != None and pre_ckpt != 'None' and pre_ckpt != '':
if 'h5' in pre_ckpt:
yolo_model_warpper.load_weights(str(pre_ckpt))
print(INFO, f' Load CKPT {str(pre_ckpt)}')
else:
print(ERROR, ' Pre CKPT path is unvalid')
# prune model
pruning_params = {
'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=initial_sparsity,
final_sparsity=final_sparsity,
begin_step=0,
end_step=h.train_epoch_step * end_epoch,
frequency=frequency)
}
if is_prune == 'True':
train_model = sparsity.prune_low_magnitude(yolo_model_warpper, **pruning_params)
else:
train_model = yolo_model_warpper
train_model.compile(
keras.optimizers.Adam(
lr=init_learning_rate,
decay=learning_rate_decay_factor),
loss=[create_loss_fn(h, obj_thresh, iou_thresh, obj_weight, noobj_weight, wh_weight, layer)
for layer in range(len(train_model.output) if isinstance(train_model.output, list) else 1)],
metrics=[Yolo_Precision(obj_thresh, name='p'), Yolo_Recall(obj_thresh, name='r')])
""" NOTE fix the dataset output shape """
shapes = (train_model.input.shape, tuple(h.output_shapes))
h.train_dataset = h.train_dataset.apply(assert_element_shape(shapes))
h.test_dataset = h.test_dataset.apply(assert_element_shape(shapes))
""" Callbacks """
if is_prune == 'True':
cbs = [
sparsity.UpdatePruningStep(),
sparsity.PruningSummaries(log_dir=str(log_dir), profile_batch=0)]
else:
cbs = [TensorBoard(str(log_dir), update_freq='batch', profile_batch=3)]
# Training
try:
train_model.fit(h.train_dataset, epochs=max_nrof_epochs,
steps_per_epoch=h.train_epoch_step, callbacks=cbs,
validation_data=h.test_dataset, validation_steps=int(h.test_epoch_step * h.validation_split))
except KeyboardInterrupt as e:
pass
if is_prune == 'True':
final_model = sparsity.strip_pruning(train_model)
prune_ckpt = log_dir / 'yolo_prune_model.h5'
keras.models.save_model(yolo_model, str(prune_ckpt), include_optimizer=False)
print()
print(INFO, f' Save Pruned Model as {str(prune_ckpt)}')
else:
keras.models.save_model(yolo_model, str(ckpt))
print()
print(INFO, f' Save Model as {str(ckpt)}')
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--train_set', type=str, help='trian file lists', default='voc')
parser.add_argument('--class_num', type=int, help='trian class num', default=20)
parser.add_argument('--pre_ckpt', type=str, help='pre-train model weights', default='None')
parser.add_argument('--model_def', type=str, help='Model definition.', default='yolo_mobilev2')
parser.add_argument('--depth_multiplier', type=float, help='mobilenet depth_multiplier', choices=[0.5, 0.75, 1.0], default=1.0)
parser.add_argument('--augmenter', type=str, help='use image augmenter', choices=['True', 'False'], default='True')
parser.add_argument('--image_size', type=int, help='net work input image size', default=(224, 320), nargs='+')
parser.add_argument('--output_size', type=int, help='net work output image size', default=(7, 10, 14, 20), nargs='+')
parser.add_argument('--batch_size', type=int, help='batch size', default=16)
parser.add_argument('--rand_seed', type=int, help='random seed', default=6)
parser.add_argument('--max_nrof_epochs', type=int, help='epoch num', default=10)
parser.add_argument('--init_learning_rate', type=float, help='init learning rate', default=0.001)
parser.add_argument('--learning_rate_decay_factor', type=float, help='learning rate decay factor', default=0)
parser.add_argument('--obj_weight', type=float, help='obj loss weight', default=5.0)
parser.add_argument('--noobj_weight', type=float, help='noobj loss weight', default=0.5)
parser.add_argument('--wh_weight', type=float, help='wh loss weight', default=0.5)
parser.add_argument('--obj_thresh', type=float, help='obj mask thresh', default=0.7)
parser.add_argument('--iou_thresh', type=float, help='iou mask thresh', default=0.3)
parser.add_argument('--vaildation_split', type=float, help='vaildation split factor', default=0.1)
parser.add_argument('--log_dir', type=str, help='log dir', default='log')
parser.add_argument('--is_prune', type=str, help='whether to prune model ', choices=['True', 'False'], default='False')
parser.add_argument('--prune_initial_sparsity', type=float, help='prune initial sparsity range = [0 ~ 1]', default=0.5)
parser.add_argument('--prune_final_sparsity', type=float, help='prune final sparsity range = [0 ~ 1]', default=0.9)
parser.add_argument('--prune_end_epoch', type=int, help='prune epochs NOTE: must < train epochs', default=5)
parser.add_argument('--prune_frequency', type=int, help='how many steps for prune once', default=100)
args = parser.parse_args(sys.argv[1:])
main(args, args.train_set, args.class_num, args.pre_ckpt,
args.model_def, args.depth_multiplier, args.augmenter,
args.image_size, args.output_size, args.batch_size, args.rand_seed, args.max_nrof_epochs,
args.init_learning_rate, args.learning_rate_decay_factor,
args.obj_weight, args.noobj_weight, args.wh_weight, args.obj_thresh, args.iou_thresh, args.vaildation_split,
args.log_dir,
args.is_prune,
args.prune_initial_sparsity,
args.prune_final_sparsity,
args.prune_end_epoch,
args.prune_frequency)