Imagenet图像分类训练总结(基于Tensorflow 2.0实现)

最近看到AWS在18年年底的一篇论文(Bag of Tricks for Image Classification with Convolutional Neural Networks),是李沐和他的同事们总结的在图像分类中用到的一些技巧,可以提高分类的准确率,我也照着论文提到的技巧测试了一下,基于Tensorflow 2.1版本,搭建了一个Darknet53的模型(这也是大名鼎鼎的YOLOV3的骨干网络),在这个基础上来对Imagenent进行分类的训练。

网络模型的搭建

首先是Darknet53网络的搭建,具体的网络结构可以参考https://github.com/pjreddie/darknet里面CFG目录下的darknet53.cfg文件。代码如下:

import tensorflow as tf
from tensorflow.keras import Model

l=tf.keras.layers
category_num = 80
vector_size = 3*(1+4+category_num)

def _conv(inputs, filters, kernel_size, strides, padding, bias=False, normalize=True, activation='relu'):
    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:
        output = l.BatchNormalization(axis=1)(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')
    output2 = _conv(output1, out_channels, 3, 1, 1, False, True, 'leaky_relu')
    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

在以上的代码中,Darknet53模型有4个输出,其中route1, route2, route3这三个是留待以后搭建YOLO V3网络时用的,在图像分类中暂时用不上。

图像预处理

论文介绍了以下的图像预处理的步骤:

  1. 随机采样图片,解码为[0,255]的32位浮点数
  2. 在图片中随机剪切一个长宽比在[3/4, 4/3]之间的矩形,其面积为图片面积的[8%, 100%]之间的一个随机数。然后把剪切后的图片缩放到224*224
  3. 随机翻转图片
  4. 随机调整图片的hue, 饱和度,明亮度,调整系数是在[0.6, 1.4]之间的一个随机值。
  5. 给图片添加PCA噪音,其系数为高斯分布(0,0.1)的一个随机值
  6. 标准化图片的RGB的值,给RGB 3个Channel分别减去123.68,116.779,103.939,然后再除以58.393,57.12,57.375 

我也遵照以上的步骤进行处理,只是在第3步翻转图片之后,我参照Darknet里面的方式,增加了一个随机旋转图片的步骤,旋转角度是在[-7, 7]之间的一个随机数。另外对于第5步的操作,在论文里面没有给出详细的介绍,我是参考mxnet里面的代码来实现的。对于图像的验证集数据来说,需要把以上的第2步改为把图片的最短边缩放为256并保持长宽比,然后在图片中间剪切一个224*224的矩形。之后跳过第3,4,5步,执行第6步即可。以下是在Tensorflow 2.1版本下的代码,构建训练集和测试集。这里我用到的Imagenent的数据是先整理为TFRECORD的格式,具体做法可以参见我之前的博客https://blog.csdn.net/gzroy/article/details/85954329

模型训练

论文中提到了以下一些技巧:

  1. 大的Batch Size的训练,随着Batch Size的增大线性增大学习率,例如Batch 128的学习率为0.1,那么Batch 256的学习率为0.2。我的显卡在FP16精度下最大只能支持128的Batch Size,因此这条技巧对我没有用。
  2. 学习率的热身,在模型刚开始训练的时候,需要从0开始逐渐增大学习率。例如设定初始学习率为0.1,在头1000个Batch的训练时,学习率是从0线性增长到0.1,这样可以有助于尽快帮助模型进入稳定学习的状态。
  3. 对残差网络的每个残差块的最后一个Batch Normalization层的γ初始化为0,这可以帮助模型在初始阶段的训练。
  4. 对权重参数的偏差项不要做L2正则化
  5. 如果显卡支持混合精度计算,则可以提高模型训练速度,同时不会网络性能不会有下降。
  6. 学习率采用余弦下降的方式,实际应用中我还是采用了Step Decay的方式,因为这样可以更可控和减少训练时间。因为余弦下降的方式学习率改变的太慢了,用Step Decay可以根据Loss值的情况来灵活调整学习率,能更快一些完成训练。当然如果显卡性能足够强大的话,用余弦下降的方式就最方便省心了。
  7. 采用标签平滑的方式来处理训练数据,例如Imagenet里面有1000个种类的图像,对于每一个特定的图像,其对应的类别的Target不是设为1,而是设为0.9,其他999个类别设置为0.1/999
  8. 知识蒸馏,就是用一个更复杂也更高准确度的教师网络,来帮助现有网络提升性能。例如用一个RESNET152的网络来帮助一个RESNET52的网络。这里我没有采用这个技巧。
  9. Mix-Up训练,也就是每次对2张采样图片进行线性整合,相应的Label也要做线性整合。这种方式需要增加训练的次数。我这里也没有采用这个技巧。不过这个技巧对于做目标识别会比较有用,可以增强模型的健壮性。

代码

完整的训练代码如下:

import tensorflow as tf
import tensorflow_addons as tfa
import math
import os
import random
import time
import numpy as np
from darknet53_model import darknet53_base
from tensorflow.keras.mixed_precision import experimental as mixed_precision
l = tf.keras.layers 

policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)

imageWidth = 224
imageHeight = 224
imageDepth = 3
batch_size = 128
resize_min = 256
train_images = 1280000
batches_per_epoch = train_images//batch_size
train_epochs = 80
total_steps = batches_per_epoch*train_epochs

random_min_aspect = 0.75
random_max_aspect = 1/0.75
random_min_area = 0.08
random_angle = 7.

initial_warmup_steps = 1000
initial_lr = 0.02

eigvec = tf.constant([[-0.5675, 0.7192, 0.4009], [-0.5808, -0.0045, -0.8140], [-0.5836, -0.6948, 0.4203]], shape=[3,3], dtype=tf.float32)
eigval = tf.constant([55.46, 4.794, 1.148], shape=[3,1], dtype=tf.float32)

mean_RGB = tf.constant([123.68, 116.779, 109.939], dtype=tf.float32)
std_RGB = tf.constant([58.393, 57.12, 57.375], dtype=tf.float32)
 
train_files_names = os.listdir('../train_tf/')
train_files = ['../train_tf/'+item for item in train_files_names]
valid_files_names = os.listdir('../valid_tf/')
valid_files = ['../valid_tf/'+item for item in valid_files_names]

# Parse TFRECORD and distort the image for train
def _parse_function(example_proto):
    features = {
        "image": tf.io.FixedLenFeature([], tf.string, default_value=""),
        "height": tf.io.FixedLenFeature([1], tf.int64, default_value=[0]),
        "width": tf.io.FixedLenFeature([1], tf.int64, default_value=[0]),
        "channels": tf.io.FixedLenFeature([1], tf.int64, default_value=[3]),
        "colorspace": tf.io.FixedLenFeature([], tf.string, default_value=""),
        "img_format": tf.io.FixedLenFeature([], tf.string, default_value=""),
        "label": tf.io.FixedLenFeature([1], tf.int64, default_value=[0]),
        "bbox_xmin": tf.io.VarLenFeature(tf.float32),
        "bbox_xmax": tf.io.VarLenFeature(tf.float32),
        "bbox_ymin": tf.io.VarLenFeature(tf.float32),
        "bbox_ymax": tf.io.VarLenFeature(tf.float32),
        "text": tf.io.FixedLenFeature([], tf.string, default_value=""),
        "filename": tf.io.FixedLenFeature([], tf.string, default_value="")
    }
    parsed_features = tf.io.parse_single_example(example_proto, features)
    image_decoded = tf.image.decode_jpeg(parsed_features["image"], channels=3)
    image_decoded = tf.cast(image_decoded, dtype=tf.float32)
    # Random crop the image 
    shape = tf.shape(image_decoded)
    height, width = shape[0], shape[1]
    random_aspect = tf.random.uniform(shape=[], minval=random_min_aspect, maxval=random_max_aspect)
    random_area = tf.random.uniform(shape=[], minval=random_min_area, maxval=1.0)
    crop_width = tf.math.sqrt(
        tf.divide(
            tf.multiply(
                tf.cast(tf.multiply(height,width), tf.float32),
                random_area),
            random_aspect)
        )
    crop_height = tf.cast(crop_width * random_aspect, tf.int32)
    crop_height = tf.cond(crop_height= initial_warmup_steps:
                tf.keras.backend.set_value(self.model.optimizer.lr, learning_rate_fn(step))
            print("Steps:{}, LR:{:6.4f}, Loss:{:4.2f}, Time:{:4.1f}s"\
                  .format(step, lr, logs['loss'], elasp_time))
    def on_epoch_end(self, epoch, logs=None):
        epoch_elasp_time = time.time()-self.epoch_starttime
        print("Epoch:{}, Top-1 Accuracy:{:5.3f}, Top-5 Accuracy:{:5.3f}, Time:{:5.1f}s"\
              .format(epoch, logs['val_output_top_1_accuracy'], logs['val_output_top_5_accuracy'], epoch_elasp_time))
    def on_epoch_begin(self, epoch, logs=None):
        tf.keras.backend.set_learning_phase(True)
        self.epoch_starttime=time.time()
    def on_test_begin(self, logs=None):
        tf.keras.backend.set_learning_phase(False)
 
tensorboard_cbk = tf.keras.callbacks.TensorBoard(log_dir='darknet53_20200203/logs')
checkpoint_cbk = tf.keras.callbacks.ModelCheckpoint(filepath='darknet53_20200203/epoch_{epoch}.h5', verbose=1)

model = darknet53_base()
model.compile(
    loss={
        'output': 
            tf.keras.losses.CategoricalCrossentropy(
                from_logits=True, label_smoothing=0.1)
    },
    optimizer=tf.keras.optimizers.SGD(
        learning_rate=0.001, momentum=0.9),
    metrics={
        'output':[
            tf.keras.metrics.CategoricalAccuracy(
                name='top_1_accuracy'),
            tf.keras.metrics.TopKCategoricalAccuracy(
                k=5, 
                name='top_5_accuracy')]
    }
)

train_data = train_input_fn()
val_data = val_input_fn()
_ = model.fit(
    train_data,
    validation_data=val_data,
    epochs=2,
    initial_epoch=0,
    verbose=0,
    callbacks=[LRCallback(time.time()), tensorboard_cbk, checkpoint_cbk],
    steps_per_epoch=5000)

总结

最终在训练了300000个Batch(30个Epoch)之后,在验证集达到了Top1 71.5%,Top5 90.6%的准确率。这个离论文提到的性能以及YOLO3的性能还有一定的差距,不过暂时已经想不到能进一步提高的方法了。

你可能感兴趣的:(机器学习,人工智能)