forked from bubbliiiing/yolov4-tiny-tf2
-
Notifications
You must be signed in to change notification settings - Fork 0
/
train_eager.py
377 lines (341 loc) · 17 KB
/
train_eager.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
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
import numpy as np
import tensorflow as tf
import tensorflow.keras.backend as K
from tensorflow.keras.layers import Input, Lambda
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import TensorBoard, ReduceLROnPlateau, EarlyStopping
from nets.yolo4_tiny import yolo_body
from nets.loss import yolo_loss
from utils.utils import get_random_data, get_random_data_with_Mosaic, rand, WarmUpCosineDecayScheduler, ModelCheckpoint
from functools import partial
from tqdm import tqdm
import time
import os
#---------------------------------------------------#
# 获得类和先验框
#---------------------------------------------------#
def get_classes(classes_path):
'''loads the classes'''
with open(classes_path) as f:
class_names = f.readlines()
class_names = [c.strip() for c in class_names]
return class_names
def get_anchors(anchors_path):
'''loads the anchors from a file'''
with open(anchors_path) as f:
anchors = f.readline()
anchors = [float(x) for x in anchors.split(',')]
return np.array(anchors).reshape(-1, 2)
#---------------------------------------------------#
# 训练数据生成器
#---------------------------------------------------#
def data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes, mosaic=False):
'''data generator for fit_generator'''
n = len(annotation_lines)
i = 0
flag = True
while True:
image_data = []
box_data = []
for b in range(batch_size):
if i==0:
np.random.shuffle(annotation_lines)
if mosaic:
if flag and (i+4) < n:
image, box = get_random_data_with_Mosaic(annotation_lines[i:i+4], input_shape)
i = (i+1) % n
else:
image, box = get_random_data(annotation_lines[i], input_shape)
i = (i+1) % n
flag = bool(1-flag)
else:
image, box = get_random_data(annotation_lines[i], input_shape)
i = (i+1) % n
image_data.append(image)
box_data.append(box)
image_data = np.array(image_data)
box_data = np.array(box_data)
y_true = preprocess_true_boxes(box_data, input_shape, anchors, num_classes)
yield image_data, y_true[0], y_true[1]
#---------------------------------------------------#
# 读入xml文件,并输出y_true
#---------------------------------------------------#
def preprocess_true_boxes(true_boxes, input_shape, anchors, num_classes):
assert (true_boxes[..., 4]<num_classes).all(), 'class id must be less than num_classes'
# 一共有三个特征层数
num_layers = len(anchors)//3
# 先验框
# 678为 142,110, 192,243, 459,401
# 345为 36,75, 76,55, 72,146
# 012为 12,16, 19,36, 40,28
anchor_mask = [[6,7,8], [3,4,5], [0,1,2]] if num_layers==3 else [[3,4,5], [1,2,3]]
true_boxes = np.array(true_boxes, dtype='float32')
input_shape = np.array(input_shape, dtype='int32') # 416,416
# 读出xy轴,读出长宽
# 中心点(m,n,2)
boxes_xy = (true_boxes[..., 0:2] + true_boxes[..., 2:4]) // 2
boxes_wh = true_boxes[..., 2:4] - true_boxes[..., 0:2]
# 计算比例
true_boxes[..., 0:2] = boxes_xy/input_shape[::-1]
true_boxes[..., 2:4] = boxes_wh/input_shape[::-1]
# m张图
m = true_boxes.shape[0]
# 得到网格的shape为13,13;26,26;52,52
grid_shapes = [input_shape//{0:32, 1:16, 2:8}[l] for l in range(num_layers)]
# y_true的格式为(m,13,13,3,85)(m,26,26,3,85)(m,52,52,3,85)
y_true = [np.zeros((m,grid_shapes[l][0],grid_shapes[l][1],len(anchor_mask[l]),5+num_classes),
dtype='float32') for l in range(num_layers)]
# [1,9,2]
anchors = np.expand_dims(anchors, 0)
anchor_maxes = anchors / 2.
anchor_mins = -anchor_maxes
# 长宽要大于0才有效
valid_mask = boxes_wh[..., 0]>0
for b in range(m):
# 对每一张图进行处理
wh = boxes_wh[b, valid_mask[b]]
if len(wh)==0: continue
# [n,1,2]
wh = np.expand_dims(wh, -2)
box_maxes = wh / 2.
box_mins = -box_maxes
# 计算真实框和哪个先验框最契合
intersect_mins = np.maximum(box_mins, anchor_mins)
intersect_maxes = np.minimum(box_maxes, anchor_maxes)
intersect_wh = np.maximum(intersect_maxes - intersect_mins, 0.)
intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]
box_area = wh[..., 0] * wh[..., 1]
anchor_area = anchors[..., 0] * anchors[..., 1]
iou = intersect_area / (box_area + anchor_area - intersect_area)
# 维度是(n) 感谢 消尽不死鸟 的提醒
best_anchor = np.argmax(iou, axis=-1)
for t, n in enumerate(best_anchor):
for l in range(num_layers):
if n in anchor_mask[l]:
# floor用于向下取整
i = np.floor(true_boxes[b,t,0]*grid_shapes[l][1]).astype('int32')
j = np.floor(true_boxes[b,t,1]*grid_shapes[l][0]).astype('int32')
# 找到真实框在特征层l中第b副图像对应的位置
k = anchor_mask[l].index(n)
c = true_boxes[b,t, 4].astype('int32')
y_true[l][b, j, i, k, 0:4] = true_boxes[b,t, 0:4]
y_true[l][b, j, i, k, 4] = 1
y_true[l][b, j, i, k, 5+c] = 1
return y_true
# 防止bug
def get_train_step_fn():
@tf.function
def train_step(imgs, yolo_loss, targets, net, optimizer, regularization):
with tf.GradientTape() as tape:
# 计算loss
P5_output, P4_output = net(imgs, training=True)
args = [P5_output, P4_output] + targets
loss_value = yolo_loss(args,anchors,num_classes,label_smoothing=label_smoothing)
if regularization:
# 加入正则化损失
loss_value = tf.reduce_sum(net.losses) + loss_value
grads = tape.gradient(loss_value, net.trainable_variables)
optimizer.apply_gradients(zip(grads, net.trainable_variables))
return loss_value
return train_step
def fit_one_epoch(net, yolo_loss, optimizer, epoch, epoch_size, epoch_size_val, gen, genval, Epoch, anchors,
num_classes, label_smoothing, regularization=False, train_step=None):
loss = 0
val_loss = 0
start_time = time.time()
with tqdm(total=epoch_size,desc=f'Epoch {epoch + 1}/{Epoch}',postfix=dict,mininterval=0.3) as pbar:
for iteration, batch in enumerate(gen):
if iteration>=epoch_size:
break
images, target0, target1 = batch[0], batch[1], batch[2]
targets = [target0, target1]
targets = [tf.convert_to_tensor(target) for target in targets]
loss_value = train_step(images, yolo_loss, targets, net, optimizer, regularization)
loss = loss + loss_value
waste_time = time.time() - start_time
pbar.set_postfix(**{'total_loss': float(loss) / (iteration + 1),
'lr' : optimizer._decayed_lr(tf.float32).numpy(),
'step/s' : waste_time})
pbar.update(1)
start_time = time.time()
print('Start Validation')
with tqdm(total=epoch_size_val, desc=f'Epoch {epoch + 1}/{Epoch}',postfix=dict,mininterval=0.3) as pbar:
for iteration, batch in enumerate(genval):
if iteration>=epoch_size_val:
break
# 计算验证集loss
images, target0, target1 = batch[0], batch[1], batch[2]
targets = [target0, target1]
targets = [tf.convert_to_tensor(target) for target in targets]
P5_output, P4_output = net(images)
args = [P5_output, P4_output] + targets
loss_value = yolo_loss(args,anchors,num_classes,label_smoothing=label_smoothing)
if regularization:
# 加入正则化损失
loss_value = tf.reduce_sum(net.losses) + loss_value
# 更新验证集loss
val_loss = val_loss + loss_value
pbar.set_postfix(**{'total_loss': float(val_loss)/ (iteration + 1)})
pbar.update(1)
print('Finish Validation')
print('Epoch:'+ str(epoch+1) + '/' + str(Epoch))
print('Total Loss: %.4f || Val Loss: %.4f ' % (loss/(epoch_size+1),val_loss/(epoch_size_val+1)))
net.save_weights('logs/Epoch%d-Total_Loss%.4f-Val_Loss%.4f.h5'%((epoch+1),loss/(epoch_size+1),val_loss/(epoch_size_val+1)))
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
#----------------------------------------------------#
# 检测精度mAP和pr曲线计算参考视频
# https://www.bilibili.com/video/BV1zE411u7Vw
#----------------------------------------------------#
if __name__ == "__main__":
#----------------------------------------------#
# 视频中说的速度慢问题已经解决了很多
# 现在train.py和train_eager.py速度差距不大
# 如果还有改进速度的地方可以私信告诉我!
#----------------------------------------------#
# 标签的位置
annotation_path = '2007_train.txt'
# 获取classes和anchor的位置
classes_path = 'model_data/voc_classes.txt'
anchors_path = 'model_data/yolo_anchors.txt'
# 预训练模型的位置
weights_path = 'model_data/yolov4_tiny_weights_coco.h5'
# 获得classes和anchor
class_names = get_classes(classes_path)
anchors = get_anchors(anchors_path)
# 一共有多少类
num_classes = len(class_names)
num_anchors = len(anchors)
#----------------------------------------------#
# 输入的shape大小
# 显存比较小可以使用416x416
# 现存比较大可以使用608x608
#----------------------------------------------#
input_shape = (416,416)
#-------------------------------#
# tricks的使用设置
#-------------------------------#
mosaic = False
Cosine_scheduler = False
label_smoothing = 0
# 是否使用正则化
regularization = True
#-------------------------------#
# Dataloder的使用
#-------------------------------#
Use_Data_Loader = True
# 输入的图像为
image_input = Input(shape=(None, None, 3))
h, w = input_shape
# 创建yolo模型
print('Create YOLOv4 model with {} anchors and {} classes.'.format(num_anchors, num_classes))
model_body = yolo_body(image_input, num_anchors//2, num_classes)
#-------------------------------------------#
# 权值文件的下载请看README
#-------------------------------------------#
print('Load weights {}.'.format(weights_path))
model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)
# 0.1用于验证,0.9用于训练
val_split = 0.1
with open(annotation_path) as f:
lines = f.readlines()
np.random.seed(10101)
np.random.shuffle(lines)
np.random.seed(None)
num_val = int(len(lines)*val_split)
num_train = len(lines) - num_val
freeze_layers = 60
for i in range(freeze_layers): model_body.layers[i].trainable = False
print('Freeze the first {} layers of total {} layers.'.format(freeze_layers, len(model_body.layers)))
#------------------------------------------------------#
# 主干特征提取网络特征通用,冻结训练可以加快训练速度
# 也可以在训练初期防止权值被破坏。
# Init_Epoch为起始世代
# Freeze_Epoch为冻结训练的世代
# Epoch总训练世代
# 提示OOM或者显存不足请调小Batch_size
#------------------------------------------------------#
if True:
Init_epoch = 0
Freeze_epoch = 50
# batch_size大小,每次喂入多少数据
batch_size = 16
# 最大学习率
learning_rate_base = 1e-3
if Use_Data_Loader:
gen = partial(data_generator, annotation_lines = lines[:num_train], batch_size = batch_size, input_shape = input_shape,
anchors = anchors, num_classes = num_classes, mosaic=mosaic)
gen = tf.data.Dataset.from_generator(gen, (tf.float32, tf.float32, tf.float32))
gen_val = partial(data_generator, annotation_lines = lines[num_train:], batch_size = batch_size,
input_shape = input_shape, anchors = anchors, num_classes = num_classes, mosaic=False)
gen_val = tf.data.Dataset.from_generator(gen_val, (tf.float32, tf.float32, tf.float32))
gen = gen.shuffle(buffer_size=batch_size).prefetch(buffer_size=batch_size)
gen_val = gen_val.shuffle(buffer_size=batch_size).prefetch(buffer_size=batch_size)
else:
gen = data_generator(lines[:num_train], batch_size, input_shape, anchors, num_classes, mosaic=mosaic)
gen_val = data_generator(lines[num_train:], batch_size, input_shape, anchors, num_classes, mosaic=False)
epoch_size = num_train//batch_size
epoch_size_val = num_val//batch_size
if Cosine_scheduler:
lr_schedule = tf.keras.experimental.CosineDecayRestarts(
initial_learning_rate = learning_rate_base,
first_decay_steps = 5*epoch_size,
t_mul = 1.0,
alpha = 1e-2
)
else:
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
initial_learning_rate=learning_rate_base,
decay_steps=epoch_size,
decay_rate=0.95,
staircase=True
)
print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))
optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)
for epoch in range(Init_epoch,Freeze_epoch):
fit_one_epoch(model_body, yolo_loss, optimizer, epoch, epoch_size, epoch_size_val,gen, gen_val,
Freeze_epoch, anchors, num_classes, label_smoothing, regularization, get_train_step_fn())
for i in range(freeze_layers): model_body.layers[i].trainable = True
# 解冻后训练
if True:
Freeze_epoch = 50
Epoch = 100
# batch_size大小,每次喂入多少数据
batch_size = 16
# 最大学习率
learning_rate_base = 1e-4
if Use_Data_Loader:
gen = partial(data_generator, annotation_lines = lines[:num_train], batch_size = batch_size, input_shape = input_shape,
anchors = anchors, num_classes = num_classes, mosaic=mosaic)
gen = tf.data.Dataset.from_generator(gen, (tf.float32, tf.float32, tf.float32))
gen_val = partial(data_generator, annotation_lines = lines[num_train:], batch_size = batch_size,
input_shape = input_shape, anchors = anchors, num_classes = num_classes, mosaic=False)
gen_val = tf.data.Dataset.from_generator(gen_val, (tf.float32, tf.float32, tf.float32))
gen = gen.shuffle(buffer_size=batch_size).prefetch(buffer_size=batch_size)
gen_val = gen_val.shuffle(buffer_size=batch_size).prefetch(buffer_size=batch_size)
else:
gen = data_generator(lines[:num_train], batch_size, input_shape, anchors, num_classes, mosaic=mosaic)
gen_val = data_generator(lines[num_train:], batch_size, input_shape, anchors, num_classes, mosaic=False)
epoch_size = num_train//batch_size
epoch_size_val = num_val//batch_size
if Cosine_scheduler:
lr_schedule = tf.keras.experimental.CosineDecayRestarts(
initial_learning_rate = learning_rate_base,
first_decay_steps = 5*epoch_size,
t_mul = 1.0,
alpha = 1e-2
)
else:
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
initial_learning_rate=learning_rate_base,
decay_steps = epoch_size,
decay_rate=0.95,
staircase=True
)
print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))
optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)
for epoch in range(Freeze_epoch,Epoch):
fit_one_epoch(model_body, yolo_loss, optimizer, epoch, epoch_size, epoch_size_val,gen, gen_val,
Epoch, anchors, num_classes, label_smoothing, regularization, get_train_step_fn())