tensorflow2.0训练网络的问题（包括BatchNormalization以及Dropout）

近期准备做一些关于深度学习图像篇的教程，主要包括分类网络，目标检测网络、图像分割网络，并以pytorch1.3以及tensorflow2.0分别去搭建实现。近期使用tensorflow2.0训练网络时遇到了很多问题，在这简单做个总结。

使用环境：Python3.6（Anaconda管理）、Tensorflow2.0.0rc1

 

1.到底使用subclassed API还是使用官方推荐的keras functional API

具体使用哪种方法，主要还是看使用者的习惯，看你喜欢使用像pytorch那样的subclassed搭建模型方式，还是喜欢使用tensorflow原来的搭建模型方式，不过根据官方教程来看，官方还是更推荐使用keras functional API的方式。如果你使用subclassed模式，这里有个坑，实例化后需要手动build一下，否则无法使用summary函数查看网络结构，最坑的是使用迁移学习时如果不提前手动build是无法载入预训练权重的，并且程序也不会报错，小心踩坑。个人建议还是使用keras functional API，坑要少点。

subclassed搭建模型方式（搭建alexNet为例）：

from tensorflow.keras import layers, models, Model, Sequential
class AlexNet(Model):
    def __init__(self, class_num=1000):
        super(AlexNet, self).__init__()
        self.features = Sequential([                                                # input(None, 224, 224, 3)
            layers.ZeroPadding2D(((1, 2), (1, 2))),                                 # output(None, 227, 227, 3)
            layers.Conv2D(48, kernel_size=11, strides=4, activation="relu"),        # output(None, 55, 55, 48)
            layers.MaxPool2D(pool_size=3, strides=2),                               # output(None, 27, 27, 48)
            layers.Conv2D(128, kernel_size=5, padding="same", activation="relu"),   # output(None, 27, 27, 128)
            layers.MaxPool2D(pool_size=3, strides=2),                               # output(None, 13, 13, 128)
            layers.Conv2D(192, kernel_size=3, padding="same", activation="relu"),   # output(None, 13, 13, 192)
            layers.Conv2D(192, kernel_size=3, padding="same", activation="relu"),   # output(None, 13, 13, 192)
            layers.Conv2D(128, kernel_size=3, padding="same", activation="relu"),   # output(None, 13, 13, 128)
            layers.MaxPool2D(pool_size=3, strides=2)])                              # output(None, 6, 6, 128)

        self.flatten = layers.Flatten()
        self.classifier = Sequential([
            layers.Dropout(0.2),
            layers.Dense(2048, activation="relu"),                                  # output(None, 2048)
            layers.Dropout(0.2),
            layers.Dense(2048, activation="relu"),                                   # output(None, 2048)
            layers.Dense(class_num),                                                # output(None, 5)
            layers.Softmax()
        ])

    def call(self, inputs):
        x = self.features(inputs)
        x = self.flatten(x)
        x = self.classifier(x)
        return x

 

keras functional API搭建模型方式（搭建alexNet为例）：

from tensorflow.keras import layers, models, Model, Sequential
def AlexNet(im_height=224, im_width=224, class_num=1000):
    # tensorflow中的tensor通道排序是NHWC                                          
    input_image = layers.Input(shape=(im_height, im_width, 3), dtype="float32")  # output(None, 224, 224, 3)
    x = layers.ZeroPadding2D(((1, 2), (1, 2)))(input_image)                      # output(None, 227, 227, 3)
    x = layers.Conv2D(48, kernel_size=11, strides=4, activation="relu")(x)       # output(None, 55, 55, 48)
    x = layers.MaxPool2D(pool_size=3, strides=2)(x)                              # output(None, 27, 27, 48)
    x = layers.Conv2D(128, kernel_size=5, padding="same", activation="relu")(x)  # output(None, 27, 27, 128)
    x = layers.MaxPool2D(pool_size=3, strides=2)(x)                              # output(None, 13, 13, 128)
    x = layers.Conv2D(192, kernel_size=3, padding="same", activation="relu")(x)  # output(None, 13, 13, 192)
    x = layers.Conv2D(192, kernel_size=3, padding="same", activation="relu")(x)  # output(None, 13, 13, 192)
    x = layers.Conv2D(128, kernel_size=3, padding="same", activation="relu")(x)  # output(None, 13, 13, 128)
    x = layers.MaxPool2D(pool_size=3, strides=2)(x)                              # output(None, 6, 6, 128)

    x = layers.Flatten()(x)                         # output(None, 6*6*128)
    x = layers.Dropout(0.2)(x)
    x = layers.Dense(2048, activation="relu")(x)    # output(None, 2048)
    x = layers.Dropout(0.2)(x)
    x = layers.Dense(2048, activation="relu")(x)    # output(None, 2048)
    x = layers.Dense(class_num)(x)                  # output(None, 5)
    predict = layers.Softmax()(x)

    model = models.Model(inputs=input_image, outputs=predict)
    return model

 

2.训练过程中使用tf.keras中的fit_generator方法还是使用低层api方法去训练

首先说下，在官方提供的一些教程中，很多例子都是使用fit方法训练，其实在真正应用中一般很少用fit方法而是使用fit_generator的方法，因为fit方法是一次性将所有样本载入内存，而fit_generator方法是分批载入内存的。在训练自己的图像样本时一般不可能一次性载入内存（内存不够），所以fit_generator的方法更常用。

使用tf.keras中的fit系列方法的好处就是训练更简单，对很底层的东西你都不需要去管，但缺点在于封装程度太高可控性差。

使用低层api的方法好处在于用户能够精确去控制训练的流程，缺点在于代码量多一点，而且需要你注意BN,Dropout层的状态（这些层在train模式以及validation模式下设置是不同的，这里坑比较多，使用fit系列方法可以不用去管）

除了以上差异外还发现使用低层的api在使用迁移学习的过程中速度明显更快，速度大概是fit系列方法的2倍以上（在CPU上跑是这样，GPU上还没测试，没有查看源码暂不知道原因），我个人推荐使用低层api的方法，能够自主控制训练过程，fit方法有些诡异。

常用fit_generator的形式如下(部分代码)：

model = MyNet(im_height=im_height, im_width=im_width, class_num=5)
model.summary()

# using keras high level api for training
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
              loss=tf.keras.losses.CategoricalCrossentropy(from_logits=False),
              metrics=["accuracy"])

callbacks = [tf.keras.callbacks.ModelCheckpoint(
    filepath='./save_weights/myNet_{epoch}.h5',
    save_best_only=True,
    save_weights_only=True,
    monitor='val_loss')]

model.fit_generator(generator=train_data_gen,
                    steps_per_epoch=total_train // batch_size,
                    epochs=epochs,
                    validation_data=val_data_gen,
                    validation_steps=total_val // batch_size,
                    callbacks=callbacks)

其中train_data_gen和val_data_gen是生成器，用于预处理并生成训练、测试样本。可以看出使用fit系列方法很简单，但你无法精准控制训练的每个batch。

 

使用低层api的训练方法如下（部分代码）：

# using keras low level api for training
loss_object = tf.keras.losses.CategoricalCrossentropy(from_logits=False)
optimizer = tf.keras.optimizers.Adam(learning_rate=0.0001)

train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.CategoricalAccuracy(name='train_accuracy')

test_loss = tf.keras.metrics.Mean(name='test_loss')
test_accuracy = tf.keras.metrics.CategoricalAccuracy(name='test_accuracy')


@tf.function
def train_step(images, labels):
    with tf.GradientTape() as tape:
        output = model(images, training=True)
        loss = loss_object(labels, output)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    train_loss(loss)
    train_accuracy(labels, output)


@tf.function
def test_step(images, labels):
    output = model(images, training=False)
    t_loss = loss_object(labels, output)

    test_loss(t_loss)
    test_accuracy(labels, output)


best_test_loss = float('inf')
for epoch in range(1, epochs+1):
    train_loss.reset_states()        # clear history info
    train_accuracy.reset_states()    # clear history info
    test_loss.reset_states()         # clear history info
    test_accuracy.reset_states()     # clear history info

    # train
    for step in range(total_train // batch_size):
        images, labels = next(train_data_gen)
        train_step(images, labels)

    # validate
    for step in range(total_val // batch_size):
        test_images, test_labels = next(val_data_gen)
        test_step(test_images, test_labels)

    template = 'Epoch {}, Loss: {}, Accuracy: {}, Test Loss: {}, Test Accuracy: {}'
    print(template.format(epoch,
                          train_loss.result(),
                          train_accuracy.result() * 100,
                          test_loss.result(),
                          test_accuracy.result() * 100))
    if test_loss.result() < best_test_loss:
        best_test_loss = test_loss.result()
        model.save_weights("./save_weights/myNet_{}.ckpt".format(epoch), save_format='tf')

从代码量上明显能够看到多了很多，但是能够精准控制训练流程。注意在每个epoch的开始都会对统计的loss以及accurate进行reset_status操作，如果不进行重置操作（即清除历史记录）会导致loss和accurate记录的不准即loss会偏大，accurate会偏低，这里要吐槽一下官方的例子，在官方的例子中并没有在每个epoch开始时清除历史记录，这会误导学习者。关于BN和Dropout的坑在下个问题中会讲到。

 

3.训练过程中BatchNormalization以及Dropout该如何处理

大家都知道Dropout和BN是解决过拟合非常常用的方法，特别是BN使用非常普遍，但如果使用不当很容易踩坑，特别是BatchNormalization（让人又爱又恨）。为什么dropout与bn层这么特殊，简单讲下原理。

dropout层常用在最后的全连接层之间，在正向传播过程中会随机以一定概率将部分节点的值置零，这样能减轻过拟合的情况。

BN层常被用在conv卷积层与relu激活层之间，用来调整正向传播过程中卷积后特征层值的分布，可防止梯度消失或爆炸以减轻过拟合的情况。

bn相比dropout更加特殊，因为bn中有四个学习参数，分别是moving_mean、moving_variance、gamma、bate（注意其中有些参数是在正向传播过程中学习更新的，坑就在于此）。

在训练过程的正向传播过程中我们希望（1）dropout随机置零节点（2）bn调节值的分布并根据样本训练自身参数

在测试推理的过程中（只存在正向传播）我们希望（1）dropout不对节点做任何操作（2）bn调节值的分布，不去学习更新自身参数

如果使用fit系列操作，我们不用去管以上说的各种问题，fit操作会自动帮我们处理，但如果使用低层的API，我们需要在train模式以及validation模式下分别传入不同的参数trainning来控制dropout以及BN的状态。其实也很简单，在train模式下传入training=True参数，在validation模式下传入training=False参数即可（model.trainable参数无法控制bn层状态）：

@tf.function
def train_step(images, labels):
    with tf.GradientTape() as tape:
        output = model(images, training=True)
        loss = loss_object(labels, output)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    train_loss(loss)
    train_accuracy(labels, output)


@tf.function
def test_step(images, labels):
    output = model(images, training=False)
    t_loss = loss_object(labels, output)

    test_loss(t_loss)
    test_accuracy(labels, output)