从第一版的yolov3(http://github.com/qqwweee/keras-yolo3)在这位q神翻译出来后,在下一直跟进yolo的发展,两年前第一次迁移了q神的keras版。最近keras版的yolov4(http://github.com/Ma-Dan/keras-yolo4)也问世了。由于tf发展到了tf2+,很多模型建立过程、命名规则、文件读取方法以及keras的支持等,都做了非常大的调整,再加上该版本的代码是延续yolov3的代码,没有使用论文的很多tricks,加上历史遗留代码存在很多的不可读因素和局部地方的小bug。因此,基于以上两点考虑,在下联合一位cv从业同学完成了基于tf2版本的、用keras编写的yolov4.
请收下传送门:https://github.com/robbebluecp/tf2-yolov4
对于这版的yolov4,我们做了如下优化:
(1)数据增强。我们在之前的resize、色彩调整、旋转的基础上,增加了mixup、mosaic、任意角度旋转(不建议用任意角
度)、背景填充、pixel等数据增强策略;
(2)模型整合。对yolo整体网络结构和局部结构做详细拆分和更详细的整合。如darknet、spp、pan等;
(3)loss优化。ciou优化、loss代码优化;
(4)convert调整。tf1+和tf2+对darknet权重文件的读取,从二进制流和命名方法上都有很大不的不同,tf2+转换非常快,且
跟tf1不能兼容。tf1和tf2对darknet训练出来的权重参数转化的h5文件,是完全不一样的。所以在使用该版本的convert.py之前。务
必确保你的tensorflow>=2.0版本!!!
(5)config配置文件取代动态传参;
(6)尽可能使np和tf分离,让训练和预测在一定程度上提速;
(7)生成器兼容、数据增强模块可扩展等其他优化。
我用voc2007数据集在V100显卡上训练到160+个epochs时,loss和val_loss差不多收敛到13+,预测准确率大约在40%-85%波动。随着预训练模型的加入或者更多epoch的训练,这个值会越来越小。
随着epoch到200+时,loss又进一步收敛到接近10(忘记留图了)。
以下展示部分数据增强的效果图:
a、rotate + mixup + resize
b、rotate + pixel + 颜色微调
c、ratate + mosaic + resize + 颜色填充
c + resize + moscai +rotate
我的小伙伴说,这些增强有些过分了,尤其是mixup,甚至连人都不太分辨的处理,这些对模型训练也是有很大影响的,所以我们用了随机数控制增强频率。在Augment类中,你可以自己控制频率,可以自定义各种增强策略。数据增强还有非常的多的方法,比如随机裁剪、图像扭曲等等。增强一时爽,一直增强一直爽。增强虽好,不不要贪强奥~~
这个框架还有很多可以完善的地方,例如loss、augment、draw等等,希望各位来添砖补瓦,集思广益,哈哈哈~~欢迎各位cv、nlp爱好者戳戳戳!
以voc2007.zip数据集为例:
如果你要训练,使用如下代码:
import loss
import config
import models
from tensorflow import keras
from generator import data_generator
class_mapping = dict(enumerate(config.classes_names))
class_mapping = {class_mapping[key]: key for key in class_mapping}
model_yolo = models.YOLO(pre_train=None)()
f = open(config.label_path)
label_lines = f.readlines()
train_lines = label_lines[:-int(len(label_lines) * config.validation_split)]
valid_lines = label_lines[-int(len(label_lines) * config.validation_split):]
h, w = config.image_input_shape
y_true = [keras.layers.Input(shape=(h // config.scale_size[l], w // config.scale_size[l], config.num_anchors, config.num_classes + 5)) for l
in range(3)]
model_loss = keras.layers.Lambda(function=loss.yolo4_loss, output_shape=(1,), name='yolo_loss')([*model_yolo.output, *y_true])
tensorboard = keras.callbacks.TensorBoard()
checkpoint = keras.callbacks.ModelCheckpoint(filepath='model_train/ep{epoch:03d}-loss{loss:.3f}-valloss{val_loss:.3f}.h5',
monitor='val_loss',
save_weights_only=True,
save_best_only=True,
period=1)
reduce_lr = keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=10, verbose=1)
early_stopping = keras.callbacks.EarlyStopping(monitor='val_loss', min_delta=0, patience=20, verbose=1)
model = keras.models.Model([model_yolo.input, *y_true], model_loss)
model.compile(optimizer=keras.optimizers.Adam(1e-4), loss={'yolo_loss': lambda y_true, y_pred: y_pred})
g_train = data_generator(label_lines=train_lines,
batch_size=config.batch_size,
input_shape=config.image_input_shape,
anchors=config.anchors,
num_classes=config.num_classes)
g_valid = data_generator(label_lines=valid_lines,
batch_size=config.batch_size,
input_shape=config.image_input_shape,
anchors=config.anchors,
num_classes=config.num_classes)
print('fire!')
model.fit(g_train,
validation_data=g_valid,
steps_per_epoch=len(label_lines) // config.batch_size,
validation_steps=int(len(label_lines) * config.validation_split * 0.2),
epochs=config.epochs,
callbacks=[tensorboard, checkpoint, reduce_lr]
)
model_yolo.save_weights('model_train/model_train_final.weights')
预测,则使用:
from tools import utils_image
import config
import cv2 as cv
import numpy as np
import eval
import models
import argparse
import tensorflow as tf
devices = tf.config.experimental.list_physical_devices('GPU')
if devices:
tf.config.experimental.set_memory_growth(devices[0], True)
parser = argparse.ArgumentParser()
parser.add_argument('-m', '--model', type=str, help='input h5 model path', default='model_train/yolov4.h5')
parser.add_argument('-i', '--image', type=str, help='input image file path', default='data/000030.jpg')
args = parser.parse_args()
model_file_path = args.model
image_file_path = args.image
class_file_path = config.classes_path
anchors = config.anchors
class_names = config.classes_names
num_anchors = len(anchors)
num_classes = len(class_names)
class_mapping = dict(enumerate(class_names))
colors = utils_image.get_random_colors(len(class_names))
class_mapping = {class_mapping[key]: key for key in class_mapping}
model = models.YOLO()()
model.load_weights('model_train/model_train_final.weights')
image = cv.imread(image_file_path)
image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
new_image = utils_image.resize_image(image, config.image_input_shape)
new_image = np.array(new_image, dtype='float32')
new_image /= 255.
new_image = np.expand_dims(new_image, 0)
feats = model.predict(new_image)
boxes, scores, classes = eval.yolo_eval(feats, anchors, len(class_names), (image.shape[0], image.shape[1]))
out_boxes, out_scores, out_classes = boxes[:5], scores[:5], classes[:5]
image = utils_image.draw_rectangle(image, boxes, scores, classes, class_names, colors, mode='pillow')
image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
cv.namedWindow("img", cv.WINDOW_NORMAL)
cv.imshow('img', image)
cv.waitKey()
对于想用darknet版本来做训练的的各位,可以移步我的另一个repo:https://github.com/robbebluecp/darknet,这个repo引自:https://github.com/AlexeyAB/darknet,因为有些地方要做修改,并且向对其中的一些代码做调整,所以独自fork了一个版本。darknet做训练有一个最大好处就是:没有python那么吃gpu!!!