15
15
16
16
17
17
18
- #---------------------------------------------------#
19
- # 获得类和先验框
20
- #---------------------------------------------------#
21
18
def get_classes (classes_path ):
22
19
'''loads the classes'''
23
20
with open (classes_path ) as f :
@@ -32,9 +29,7 @@ def get_anchors(anchors_path):
32
29
anchors = [float (x ) for x in anchors .split (',' )]
33
30
return np .array (anchors ).reshape (- 1 , 2 )
34
31
35
- #---------------------------------------------------#
36
- # 训练数据生成器
37
- #---------------------------------------------------#
32
+
38
33
def data_generator (annotation_lines , batch_size , input_shape , anchors , num_classes , mosaic = False , random = True ):
39
34
n = len (annotation_lines )
40
35
i = 0
@@ -63,75 +58,49 @@ def data_generator(annotation_lines, batch_size, input_shape, anchors, num_class
63
58
y_true = preprocess_true_boxes (box_data , input_shape , anchors , num_classes )
64
59
yield [image_data , * y_true ], np .zeros (batch_size )
65
60
66
- #---------------------------------------------------#
67
- # 读入xml文件,并输出y_true
68
- #---------------------------------------------------#
61
+
69
62
def preprocess_true_boxes (true_boxes , input_shape , anchors , num_classes ):
70
63
assert (true_boxes [..., 4 ]< num_classes ).all (), 'class id must be less than num_classes'
71
- # 一共有两个特征层数
64
+
72
65
num_layers = len (anchors )// 3
73
- #-----------------------------------------------------------#
74
- # 13x13的特征层对应的anchor是[81,82], [135,169], [344,319]
75
- # 26x26的特征层对应的anchor是[23,27], [37,58], [81,82]
76
- #-----------------------------------------------------------#
66
+
77
67
anchor_mask = [[6 ,7 ,8 ], [3 ,4 ,5 ], [0 ,1 ,2 ]] if num_layers == 3 else [[3 ,4 ,5 ], [1 ,2 ,3 ]]
78
68
79
- #-----------------------------------------------------------#
80
- # 获得框的坐标和图片的大小
81
- #-----------------------------------------------------------#
69
+
82
70
true_boxes = np .array (true_boxes , dtype = 'float32' )
83
71
input_shape = np .array (input_shape , dtype = 'int32' )
84
- #-----------------------------------------------------------#
85
- # 通过计算获得真实框的中心和宽高
86
- # 中心点(m,n,2) 宽高(m,n,2)
87
- #-----------------------------------------------------------#
72
+
88
73
boxes_xy = (true_boxes [..., 0 :2 ] + true_boxes [..., 2 :4 ]) // 2
89
74
boxes_wh = true_boxes [..., 2 :4 ] - true_boxes [..., 0 :2 ]
90
- #-----------------------------------------------------------#
91
- # 将真实框归一化到小数形式
92
- #-----------------------------------------------------------#
75
+
93
76
true_boxes [..., 0 :2 ] = boxes_xy / input_shape [::- 1 ]
94
77
true_boxes [..., 2 :4 ] = boxes_wh / input_shape [::- 1 ]
95
78
96
- # m为图片数量,grid_shapes为网格的shape
79
+
97
80
m = true_boxes .shape [0 ]
98
81
grid_shapes = [input_shape // {0 :32 , 1 :16 , 2 :8 }[l ] for l in range (num_layers )]
99
- #-----------------------------------------------------------#
100
- # y_true的格式为(m,13,13,3,85)(m,26,26,3,85)(m,52,52,3,85)
101
- #-----------------------------------------------------------#
82
+
102
83
y_true = [np .zeros ((m ,grid_shapes [l ][0 ],grid_shapes [l ][1 ],len (anchor_mask [l ]),5 + num_classes ),
103
84
dtype = 'float32' ) for l in range (num_layers )]
104
85
105
- #-----------------------------------------------------------#
106
- # [6,2] -> [1,6,2]
107
- #-----------------------------------------------------------#
86
+
108
87
anchors = np .expand_dims (anchors , 0 )
109
88
anchor_maxes = anchors / 2.
110
89
anchor_mins = - anchor_maxes
111
90
112
- #-----------------------------------------------------------#
113
- # 长宽要大于0才有效
114
- #-----------------------------------------------------------#
91
+
115
92
valid_mask = boxes_wh [..., 0 ]> 0
116
93
117
94
for b in range (m ):
118
- # 对每一张图进行处理
95
+
119
96
wh = boxes_wh [b , valid_mask [b ]]
120
97
if len (wh )== 0 : continue
121
- #-----------------------------------------------------------#
122
- # [n,2] -> [n,1,2]
123
- #-----------------------------------------------------------#
98
+
124
99
wh = np .expand_dims (wh , - 2 )
125
100
box_maxes = wh / 2.
126
101
box_mins = - box_maxes
127
102
128
- #-----------------------------------------------------------#
129
- # 计算所有真实框和先验框的交并比
130
- # intersect_area [n,6]
131
- # box_area [n,1]
132
- # anchor_area [1,6]
133
- # iou [n,6]
134
- #-----------------------------------------------------------#
103
+
135
104
intersect_mins = np .maximum (box_mins , anchor_mins )
136
105
intersect_maxes = np .minimum (box_maxes , anchor_maxes )
137
106
intersect_wh = np .maximum (intersect_maxes - intersect_mins , 0. )
@@ -141,141 +110,76 @@ def preprocess_true_boxes(true_boxes, input_shape, anchors, num_classes):
141
110
anchor_area = anchors [..., 0 ] * anchors [..., 1 ]
142
111
143
112
iou = intersect_area / (box_area + anchor_area - intersect_area )
144
- #-----------------------------------------------------------#
145
- # 维度是[n,] 感谢 消尽不死鸟 的提醒
146
- #-----------------------------------------------------------#
113
+
147
114
best_anchor = np .argmax (iou , axis = - 1 )
148
115
149
116
for t , n in enumerate (best_anchor ):
150
- #-----------------------------------------------------------#
151
- # 找到每个真实框所属的特征层
152
- #-----------------------------------------------------------#
117
+
153
118
for l in range (num_layers ):
154
119
if n in anchor_mask [l ]:
155
- #-----------------------------------------------------------#
156
- # floor用于向下取整,找到真实框所属的特征层对应的x、y轴坐标
157
- #-----------------------------------------------------------#
120
+
158
121
i = np .floor (true_boxes [b ,t ,0 ] * grid_shapes [l ][1 ]).astype ('int32' )
159
122
j = np .floor (true_boxes [b ,t ,1 ] * grid_shapes [l ][0 ]).astype ('int32' )
160
- #-----------------------------------------------------------#
161
- # k指的的当前这个特征点的第k个先验框
162
- #-----------------------------------------------------------#
123
+
163
124
k = anchor_mask [l ].index (n )
164
- #-----------------------------------------------------------#
165
- # c指的是当前这个真实框的种类
166
- #-----------------------------------------------------------#
125
+
167
126
c = true_boxes [b , t , 4 ].astype ('int32' )
168
- #-----------------------------------------------------------#
169
- # y_true的shape为(m,13,13,3,85)(m,26,26,3,85)(m,52,52,3,85)
170
- # 最后的85可以拆分成4+1+80,4代表的是框的中心与宽高、
171
- # 1代表的是置信度、80代表的是种类
172
- #-----------------------------------------------------------#
127
+
173
128
y_true [l ][b , j , i , k , 0 :4 ] = true_boxes [b , t , 0 :4 ]
174
129
y_true [l ][b , j , i , k , 4 ] = 1
175
130
y_true [l ][b , j , i , k , 5 + c ] = 1
176
131
177
132
return y_true
178
133
179
- #----------------------------------------------------#
180
- # 检测精度mAP和pr曲线计算参考视频
181
- # https://www.bilibili.com/video/BV1zE411u7Vw
182
- #----------------------------------------------------#
134
+
183
135
if __name__ == "__main__" :
184
- #----------------------------------------------------#
185
- # 获得图片路径和标签
186
- #----------------------------------------------------#
136
+
187
137
annotation_path = '/content/drive/MyDrive/yolov4_tiny/yolov4-tiny-keras-master/new_train.txt'
188
- #------------------------------------------------------#
189
- # 训练后的模型保存的位置,保存在logs文件夹里面
190
- #------------------------------------------------------#
138
+
191
139
log_dir = 'logs/'
192
- #----------------------------------------------------#
193
- # classes和anchor的路径,非常重要
194
- # 训练前一定要修改classes_path,使其对应自己的数据集
195
- #----------------------------------------------------#
140
+
196
141
classes_path = '/content/drive/MyDrive/yolov4_tiny/yolov4-tiny-keras-master/voc_classes.txt'
197
142
anchors_path = '/content/drive/MyDrive/yolov4_tiny/yolov4-tiny-keras-master/yolo_anchors.txt'
198
- #------------------------------------------------------#
199
- # 权值文件请看README,百度网盘下载
200
- # 训练自己的数据集时提示维度不匹配正常
201
- # 预测的东西都不一样了自然维度不匹配
202
- #------------------------------------------------------#
143
+
203
144
weights_path = '/content/drive/MyDrive/yolov4_tiny/yolov4-tiny-keras-master/yolov4_tiny_weights_voc.h5'
204
- #------------------------------------------------------#
205
- # 训练用图片大小
206
- # 一般在416x416和608x608选择
207
- #------------------------------------------------------#
145
+
208
146
input_shape = (416 ,416 )
209
- #------------------------------------------------------#
210
- # 是否对损失进行归一化,用于改变loss的大小
211
- # 用于决定计算最终loss是除上batch_size还是除上正样本数量
212
- #------------------------------------------------------#
147
+
213
148
normalize = False
214
149
215
- #----------------------------------------------------#
216
- # 获取classes和anchor
217
- #----------------------------------------------------#
218
150
class_names = get_classes (classes_path )
219
151
anchors = get_anchors (anchors_path )
220
- #------------------------------------------------------#
221
- # 一共有多少类和多少先验框
222
- #------------------------------------------------------#
223
152
num_classes = len (class_names )
224
153
num_anchors = len (anchors )
225
- #------------------------------------------------------#
226
- # Yolov4的tricks应用
227
- # mosaic 马赛克数据增强 True or False
228
- # 实际测试时mosaic数据增强并不稳定,所以默认为False
229
- # Cosine_scheduler 余弦退火学习率 True or False
230
- # label_smoothing 标签平滑 0.01以下一般 如0.01、0.005
231
- #------------------------------------------------------#
232
154
mosaic = False
233
155
Cosine_scheduler = False
234
156
label_smoothing = 0
235
157
236
158
K .clear_session ()
237
- #------------------------------------------------------#
238
- # 创建yolo模型
239
- #------------------------------------------------------#
159
+
240
160
image_input = Input (shape = (None , None , 3 ))
241
161
h , w = input_shape
242
162
print ('Create YOLOv4-Tiny model with {} anchors and {} classes.' .format (num_anchors , num_classes ))
243
163
model_body = yolo_body (image_input , num_anchors // 2 , num_classes )
244
164
245
- #------------------------------------------------------#
246
- # 载入预训练权重
247
- #------------------------------------------------------#
165
+
248
166
print ('Load weights {}.' .format (weights_path ))
249
167
model_body .load_weights (weights_path , by_name = True , skip_mismatch = True )
250
168
251
- #------------------------------------------------------#
252
- # 在这个地方设置损失,将网络的输出结果传入loss函数
253
- # 把整个模型的输出作为loss
254
- #------------------------------------------------------#
255
169
y_true = [Input (shape = (h // {0 :32 , 1 :16 }[l ], w // {0 :32 , 1 :16 }[l ], num_anchors // 2 , num_classes + 5 )) for l in range (2 )]
256
170
loss_input = [* model_body .output , * y_true ]
257
171
model_loss = Lambda (yolo_loss , output_shape = (1 ,), name = 'yolo_loss' ,
258
172
arguments = {'anchors' : anchors , 'num_classes' : num_classes , 'ignore_thresh' : 0.5 , 'label_smoothing' : label_smoothing , 'normalize' :normalize })(loss_input )
259
173
260
174
model = Model ([model_body .input , * y_true ], model_loss )
261
175
262
- #-------------------------------------------------------------------------------#
263
- # 训练参数的设置
264
- # logging表示tensorboard的保存地址
265
- # checkpoint用于设置权值保存的细节,period用于修改多少epoch保存一次
266
- # reduce_lr用于设置学习率下降的方式
267
- # early_stopping用于设定早停,val_loss多次不下降自动结束训练,表示模型基本收敛
268
- #-------------------------------------------------------------------------------#
176
+
269
177
logging = TensorBoard (log_dir = log_dir )
270
178
checkpoint = ModelCheckpoint (log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5' ,
271
179
monitor = 'val_loss' , save_weights_only = True , save_best_only = False , period = 1 )
272
180
early_stopping = EarlyStopping (monitor = 'val_loss' , min_delta = 0 , patience = 10 , verbose = 1 )
273
181
274
- #----------------------------------------------------------------------#
275
- # 验证集的划分在train.py代码里面进行
276
- # 2007_test.txt和2007_val.txt里面没有内容是正常的。训练不会使用到。
277
- # 当前划分方式下,验证集和训练集的比例为1:9
278
- #----------------------------------------------------------------------#
182
+
279
183
val_split = 0.1
280
184
with open (annotation_path ) as f :
281
185
lines = f .readlines ()
@@ -285,33 +189,26 @@ def preprocess_true_boxes(true_boxes, input_shape, anchors, num_classes):
285
189
num_val = int (len (lines )* val_split )
286
190
num_train = len (lines ) - num_val
287
191
288
- #------------------------------------------------------#
289
- # 主干特征提取网络特征通用,冻结训练可以加快训练速度
290
- # 也可以在训练初期防止权值被破坏。
291
- # Init_Epoch为起始世代
292
- # Freeze_Epoch为冻结训练的世代
293
- # Epoch总训练世代
294
- # 提示OOM或者显存不足请调小Batch_size
295
- #------------------------------------------------------#
192
+
296
193
freeze_layers = 60
297
194
for i in range (freeze_layers ): model_body .layers [i ].trainable = False
298
195
print ('Freeze the first {} layers of total {} layers.' .format (freeze_layers , len (model_body .layers )))
299
196
300
- # 调整非主干模型first
197
+
301
198
if True :
302
199
Init_epoch = 0
303
200
Freeze_epoch = 50
304
201
batch_size = 32
305
202
learning_rate_base = 1e-3
306
203
307
204
if Cosine_scheduler :
308
- # 预热期
205
+
309
206
warmup_epoch = int ((Freeze_epoch - Init_epoch )* 0.2 )
310
- # 总共的步长
207
+
311
208
total_steps = int ((Freeze_epoch - Init_epoch ) * num_train / batch_size )
312
- # 预热步长
209
+
313
210
warmup_steps = int (warmup_epoch * num_train / batch_size )
314
- # 学习率
211
+
315
212
reduce_lr = WarmUpCosineDecayScheduler (learning_rate_base = learning_rate_base ,
316
213
total_steps = total_steps ,
317
214
warmup_learning_rate = 1e-4 ,
@@ -336,21 +233,21 @@ def preprocess_true_boxes(true_boxes, input_shape, anchors, num_classes):
336
233
337
234
for i in range (freeze_layers ): model_body .layers [i ].trainable = True
338
235
339
- # 解冻后训练
236
+
340
237
if True :
341
238
Freeze_epoch = 50
342
239
Epoch = 100
343
240
batch_size = 16
344
241
learning_rate_base = 1e-4
345
242
346
243
if Cosine_scheduler :
347
- # 预热期
244
+
348
245
warmup_epoch = int ((Epoch - Freeze_epoch )* 0.2 )
349
- # 总共的步长
246
+
350
247
total_steps = int ((Epoch - Freeze_epoch ) * num_train / batch_size )
351
- # 预热步长
248
+
352
249
warmup_steps = int (warmup_epoch * num_train / batch_size )
353
- # 学习率
250
+
354
251
reduce_lr = WarmUpCosineDecayScheduler (learning_rate_base = learning_rate_base ,
355
252
total_steps = total_steps ,
356
253
warmup_learning_rate = 1e-5 ,
0 commit comments