YOLO V3基于Tensorflow 2.0的完整实现

原创

已于 2023-04-06 19:34:45 修改 · 9.2k 阅读

128

于 2020-06-21 23:16:25 首次发布

本文详细介绍使用TensorFlow实现YOLOv3物体检测模型的过程，包括训练数据准备、网络模型搭建、训练过程及性能评估。

如果对Tensorflow实现最新的Yolo v7算法感兴趣的朋友，可以参见我最新发布的文章，Yolo v7的最简TensorFlow实现_gzroy的博客-CSDN博客

YOLO V3版本是一个强大和快速的物体检测模型，同时原理上也相对简单。我之前的博客中已经介绍了如何用Tensorflow来实现YOLO V1版本，之后我自己也用Tensorflow 1.X版本实现了YOLO V3，现在Tensorflow演进到了2.0版本，相比较１.X版本做了很大的改进，也更加易用了，因此我记录一下如何用Tensorflow 2.0版本来实现YOLO V3。网上能找到的很多Tensorflow YOLO V3的代码都没有完整的一个训练过程，基本上都是转换和加载YOLO的作者在Darknet上发布的训练好的权重数据，直接进行检测的。我的这个代码实现了完整的训练流程，包括了搭建基础架构网络Darknet53进行Imagenet预训练，以及增加YOLO V3网络模块进行物体检测训练，对模型的训练效果进行评测，以及用训练好的模型进行物体检测的过程。

训练数据的准备

需要准备两份训练数据，一个是Imagenent的物体分类数据，包括了1000种类别的物体的数据，共128万张图片。数据集需要预先处理为TFRECORD格式，具体过程可以参见我之前的博客基于Tensorflow的Imagenet数据集的完整处理过程（包括物体标识框BBOX的处理）_valid_classes_gzroy的博客-CSDN博客。第二个训练数据是物体检测的数据，目前有很多个数据集可以采用，例如COCO数据集（包括80种物体的检测框），OpenImage，Pascal VOC等等，比较流行的是COCO数据集，大部分物体检测的论文都会基于这个数据集来提供性能指标。我也采用COCO数据集，同样也是预处理为TFRECORD格式，具体过程可以参见我的另一篇博客基于Tensorflow对COCO目标检测数据进行预处理_gzroy的博客-CSDN博客

网络模型的搭建

按照YOLO V3论文的描述，基础网络架构是一个叫做Darknet53的网络模型，共有53个卷积层，其网络架构如下：

用Tensorflow可以很方便的构建一个Darknet53模型，代码如下：

import tensorflow as tf
from tensorflow.keras import Model
l=tf.keras.layers

def _conv(inputs, filters, kernel_size, strides, padding, bias=False, normalize=True, activation='relu', last=False):
    output = inputs
    padding_str = 'same'
    if padding>0:
        output = l.ZeroPadding2D(padding=padding, data_format='channels_first')(output)
        padding_str = 'valid'
    output = l.Conv2D(filters, kernel_size, strides, padding_str, \
                  'channels_first', use_bias=bias, \
                  kernel_initializer='he_normal', \
                  kernel_regularizer=tf.keras.regularizers.l2(l=5e-4))(output)
    if normalize:
        if not last:
            output = l.BatchNormalization(axis=1)(output)
        else:
            output = l.BatchNormalization(axis=1, gamma_initializer='zeros')(output)
    if activation=='relu':
        output = l.ReLU()(output)
    if activation=='relu6':
        output = l.ReLU(max_value=6)(output)
    if activation=='leaky_relu':
        output = l.LeakyReLU(alpha=0.1)(output)
    return output

def _residual(inputs, out_channels, activation='relu', name=None):
    output1 = _conv(inputs, out_channels//2, 1, 1, 0, False, True, 'leaky_relu', False)
    output2 = _conv(output1, out_channels, 3, 1, 1, False, True, 'leaky_relu', True)
    output = l.Add(name=name)([inputs, output2])
    return output 

def darknet53_base():
    image = tf.keras.Input(shape=(3,None,None))
    net = _conv(image, 32, 3, 1, 1, False, True, 'leaky_relu')     #32*H*W
    net = _conv(net, 64, 3, 2, 1, False, True, 'leaky_relu')       #64*H/2*W/2
    net = _residual(net, 64, 'leaky_relu')                         #64*H/2*W/2
    net = _conv(net, 128, 3, 2, 1, False, True, 'leaky_relu')      #128*H/4*W/4
    net = _residual(net, 128, 'leaky_relu')                        #128*H/4*W/4
    net = _residual(net, 128, 'leaky_relu')                        #128*H/4*W/4
    net = _conv(net, 256, 3, 2, 1, False, True, 'leaky_relu')      #256*H/8*W/8
    net = _residual(net, 256, 'leaky_relu')                        #256*H/8*W/8
    net = _residual(net, 256, 'leaky_relu')                        #256*H/8*W/8
    net = _residual(net, 256, 'leaky_relu')                        #256*H/8*W/8
    net = _residual(net, 256, 'leaky_relu')                        #256*H/8*W/8
    net = _residual(net, 256, 'leaky_relu')                        #256*H/8*W/8
    net = _residual(net, 256, 'leaky_relu')                        #256*H/8*W/8
    net = _residual(net, 256, 'leaky_relu')                        #256*H/8*W/8
    net = _residual(net, 256, 'leaky_relu')                        #256*H/8*W/8
    route1 = l.Activation('linear', dtype='float32', name='route1')(net)
    net = _conv(net, 512, 3, 2, 1, False, True, 'leaky_relu')   #512*H/16*W/16
    net = _residual(net, 512, 'leaky_relu')                        #512*H/16*W/16
    net = _residual(net, 512, 'leaky_relu')                        #512*H/16*W/16
    net = _residual(net, 512, 'leaky_relu')                        #512*H/16*W/16
    net = _residual(net, 512, 'leaky_relu')                        #512*H/16*W/16
    net = _residual(net, 512, 'leaky_relu')                        #512*H/16*W/16
    net = _residual(net, 512, 'leaky_relu')                        #512*H/16*W/16
    net = _residual(net, 512, 'leaky_relu')                        #512*H/16*W/16
    net = _residual(net, 512, 'leaky_relu')                        #512*H/16*W/16
    route2 = l.Activation('linear', dtype='float32', name='route2')(net)
    net = _conv(net, 1024, 3, 2, 1, False, True, 'leaky_relu')     #1024*H/32*W/32
    net = _residual(net, 1024, 'leaky_relu')                       #1024*H/32*W/32
    net = _residual(net, 1024, 'leaky_relu')                       #1024*H/32*W/32
    net = _residual(net, 1024, 'leaky_relu')                       #1024*H/32*W/32
    net = _residual(net, 1024, 'leaky_relu')                       #1024*H/32*W/32
    route3 = l.Activation('linear', dtype='float32', name='route3')(net)
    net = tf.reduce_mean(net, axis=[2,3], keepdims=True)
    net = _conv(net, 1000, 1, 1, 0, True, False, 'linear')         #1000
    net = l.Flatten(data_format='channels_first', name='logits')(net)
    net = l.Activation('linear', dtype='float32', name='output')(net)
    model = tf.keras.Model(inputs=image, outputs=[net, route1, route2, route3])
    return model

我们需要先基于这个骨干网络架构来进行Imagenet的预训练，以提取有效的图片内容的特征数据。我用这个网络训练了30个EPOCH，最终到达Top-1 71%，Top-5 91%的准确率。具体的训练过程可以见我的博客Imagenet图像分类训练总结（基于Tensorflow 2.0实现）_keras .image_gzroy的博客-CSDN博客

训练好了骨干网络之后，我们就可以在这个网络的基础上再增加相应的卷积层，实现图像特征金字塔(FPN)的架构，这里我们会用到骨干网络输出的route1, route2, route3这几个不同图像分辨率的特征值，最终构建一个可以对图片进行下采样8倍，16倍和32倍的基于网格的检测系统，例如训练图片的分辨率为416*416，那么将输出52*52, 26*26, 13*13这三个不同维度的检测结果。具体的原理可以参见网上的一些文章，例如：我这里参照Darknet的源代码来搭建了一个YOLO V3的网络，代码如下：

category_num = 80
vector_size = 3*(1+4+category_num)
def darknet53_yolov3():
    route1 = tf.keras.Input(shape=(256,None,None), name='input1')        #256*H/8*W/8
    route2 = tf.keras.Input(shape=(512,None,None), name='input2')        #256*H/16*W/16
    route3 = tf.keras.Input(shape=(1024,None,None), name='input3')       #256*H/32*W/32
    net = _conv(route3, 512, 1, 1, 0, False, True, 'leaky_relu')         #512*H/32*W/32
    net = _conv(net, 1024, 3, 1, 1, False, True, 'leaky_relu')           #1024*H/32*W/32
    net = _conv(net, 512, 1, 1, 0, False, True, 'leaky_relu')            #512*H/32*W/32
    net = _conv(net, 1024, 3, 1, 1, False, True, 'leaky_relu')           #1024*H/32*W/32
    net = _conv(net, 512, 1, 1, 0, False, True, 'leaky_relu')            #512*H/32*W/32
    route4 = tf.identity(net, 'route4')
    net = _conv(net, 1024, 3, 1, 1, False, True, 'leaky_relu')           #1024*H/32*W/32
    predict1 = _conv(net, vector_size, 1, 1, 0, True, False, 'linear')   #vector_size*H/32*W/32
    predict1 = l.Activation('linear', dtype='float32')(predict1)
    predict1 = l.Reshape((vector_size, imageHeight//32*imageWidth//32))(predict1)
    net = _conv(route4, 256, 1, 1, 0, False, True, 'leaky_relu')         #256*H/32*W/32
    net = l.UpSampling2D((2,2),"channels_first",'nearest')(net)    #256*H/16*W/16
    net = l.Concatenate(axis=1)([route2, net])                     #768*H/16*W/16
    net = _conv(net, 256, 1, 1, 0, False, True, 'leaky_relu')            #256*H/16*W/16
    net = _conv(net, 512, 3, 1, 1, False, True, 'leaky_relu')            #512*H/16*W/16
    net = _conv(net, 256, 1, 1, 0, False, True, 'leaky_relu')            #256*H/16*W/16
    net = _conv(net, 512, 3, 1, 1, False, True, 'leaky_relu')            #512*H/16*W/16
    net = _conv(net, 256, 1, 1, 0, False, True, 'leaky_relu')            #256*H/16*W/16
    route5 = tf.identity(net, 'route5')
    net = _conv(net, 512, 3, 1, 1, False, True, 'leaky_relu')            #512*H/16*W/16
    predict2 = _conv(net, vector_size, 1, 1, 0, True, False, 'linear')   #vector_size*H/16*W/16
    predict2 = l.Activation('linear', dtype='float32')(predict2)
    predict2 = l.Reshape((vector_size, imageHeight//16*imageWidth//16))(predict2)
    net = _conv(route5, 128, 1, 1, 0, False, True, 'leaky_relu')         #128*H/16*W/16
    net = l.UpSampling2D((2,2),"channels_first",'nearest')(net)    #128*H/8*W/8
    net = l.Concatenate(axis=1)([route1, net])                     #384*H/8*W/8
    net = _conv(net, 128, 1, 1, 0, False, True, 'leaky_relu')            #128*H/8*W/8
    net = _conv(net, 256, 3, 1, 1, False, True, 'leaky_relu')            #256*H/8*W/8
    net = _conv(net, 128, 1, 1, 0, False, True, 'leaky_relu')            #128*H/8*W/8
    net = _conv(net, 256, 3, 1, 1, False, True, 'leaky_relu')            #256*H/8*W/8

标签