-
Notifications
You must be signed in to change notification settings - Fork 23
/
train.py
executable file
·163 lines (130 loc) · 6.16 KB
/
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
154
155
156
157
158
159
160
161
162
163
# -*- coding: utf-8 -*-
# @Time : 2020/3/24
# @File : train.py
# @Software: PyCharm
import os
import sys
import time
import tensorflow as tf
from absl import flags, logging, app
from absl.flags import FLAGS
from components import config
from components.lr_scheduler import MultiStepWarmUpLR
from components.prior_box import priors_box
from components.utils import set_memory_growth
from dataset.tf_dataset_preprocess import load_dataset
from network.losses import MultiBoxLoss
from network.network import SlimModel
flags.DEFINE_string('gpu', '0', 'which gpu to use')
def main(_):
global load_t1
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu # CPU:'-1'
logger = tf.get_logger()
logger.disabled = True
logger.setLevel(logging.FATAL)
set_memory_growth()
weights_dir = 'checkpoints/'
if not os.path.exists(weights_dir):
os.mkdir(weights_dir)
# if os.path.exists('logs'):
# shutil.rmtree('logs')
logging.info("Load configuration...")
cfg = config.cfg
label_classes = cfg['labels_list']
logging.info(f"Total image sample:{cfg['dataset_len']},Total classes number:"
f"{len(label_classes)},classes list:{label_classes}")
logging.info("Compute priors boxes...")
priors, num_cell = priors_box(cfg)
logging.info(f"Prior boxes number:{len(priors)},default anchor box number per feature map cell:{num_cell}")
logging.info("Loading dataset...")
train_dataset = load_dataset(cfg, priors, shuffle=True, train=True)
# val_dataset = load_dataset(cfg, priors, shuffle=False, train=False)
logging.info("Create Model...")
try:
model = SlimModel(cfg=cfg, num_cell=num_cell, training=True)
model.summary()
tf.keras.utils.plot_model(model, to_file=os.path.join(os.getcwd(), 'model.png'),
show_shapes=True, show_layer_names=True)
except Exception as e:
logging.error(e)
logging.info("Create network failed.")
sys.exit()
if cfg['resume']:
# Training from latest weights
paths = [os.path.join(weights_dir, path)
for path in os.listdir(weights_dir)]
latest = sorted(paths, key=os.path.getmtime)[-1]
model.load_weights(latest)
init_epoch = int(os.path.splitext(latest)[0][-3:])
else:
# Training from scratch
init_epoch = -1
steps_per_epoch = cfg['dataset_len'] // cfg['batch_size']
# val_steps_per_epoch = cfg['val_len'] // cfg['batch_size']
logging.info(f"steps_per_epoch:{steps_per_epoch}")
logging.info("Define optimizer and loss computation and so on...")
# learning_rate =tf.keras.optimizers.schedules.ExponentialDecay(initial_learning_rate=1e-3,
# decay_steps=20000,
# decay_rate=0.96)
# optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
learning_rate = MultiStepWarmUpLR(
initial_learning_rate=cfg['init_lr'],
lr_steps=[e * steps_per_epoch for e in cfg['lr_decay_epoch']],
lr_rate=cfg['lr_rate'],
warmup_steps=cfg['warmup_epoch'] * steps_per_epoch,
min_lr=cfg['min_lr'])
optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=cfg['momentum'], nesterov=True)
multi_loss = MultiBoxLoss(num_class=len(label_classes), neg_pos_ratio=3)
train_log_dir = 'logs/train'
train_summary_writer = tf.summary.create_file_writer(train_log_dir)
@tf.function
def train_step(inputs, labels):
with tf.GradientTape() as tape:
predictions = model(inputs, training=True)
losses = {}
losses['reg'] = tf.reduce_sum(model.losses) # unused. Init for redefine network
losses['loc'], losses['class'] = multi_loss(labels, predictions)
total_loss = tf.add_n([l for l in losses.values()])
grads = tape.gradient(total_loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
return total_loss, losses
for epoch in range(init_epoch + 1, cfg['epoch']):
try:
start = time.time()
avg_loss = 0.0
for step, (inputs, labels) in enumerate(train_dataset.take(steps_per_epoch)):
load_t0 = time.time()
total_loss, losses = train_step(inputs, labels)
avg_loss = (avg_loss * step + total_loss.numpy()) / (step + 1)
load_t1 = time.time()
batch_time = load_t1 - load_t0
steps = steps_per_epoch * epoch + step
with train_summary_writer.as_default():
tf.summary.scalar('loss/total_loss', total_loss, step=steps)
for k, l in losses.items():
tf.summary.scalar('loss/{}'.format(k), l, step=steps)
tf.summary.scalar('learning_rate', optimizer.lr(steps), step=steps)
print(
f"\rEpoch: {epoch + 1}/{cfg['epoch']} | Batch {step + 1}/{steps_per_epoch} | Batch time {batch_time:.3f} || Loss: {total_loss:.6f} | loc loss:{losses['loc']:.6f} | class loss:{losses['class']:.6f} ",
end='', flush=True)
print(
f"\nEpoch: {epoch + 1}/{cfg['epoch']} | Epoch time {(load_t1 - start):.3f} || Average Loss: {avg_loss:.6f}")
with train_summary_writer.as_default():
tf.summary.scalar('loss/avg_loss', avg_loss, step=epoch)
if (epoch + 1) % cfg['save_freq'] == 0:
filepath = os.path.join(weights_dir, f'weights_epoch_{(epoch + 1):03d}.h5')
model.save_weights(filepath)
if os.path.exists(filepath):
print(f">>>>>>>>>>Save weights file at {filepath}<<<<<<<<<<")
except KeyboardInterrupt:
print('interrupted')
# filepath = os.path.join(weights_dir, 'weights_last.h5')
# model.save_weights(filepath)
# print(f'model saved into: {filepath}')
exit(0)
if __name__ == '__main__':
try:
app.run(main)
except SystemExit:
pass