保姆级 Keras 实现 Faster R-CNN 八

上一篇文章 我们完成了 RPN 网络的训练, 还提到了按照文章中的步骤训练时 Loss 可能会卡住, 不会得到一个好的模型. 或者 Loss 不会卡住, 但是下降得特别慢. 这篇文章就来讲如果破解

一. 预训练模型

用 VOC2007 训练 Faster R-CNN 不太好训练主要是数据分布很不 “集中”, 再加上又是多输出, 每个输出的梯度都会相互影响. 一个要往东走, 一个要向西, 很难达成一致. 要是运气好, 初始化的参数使多个输出都向着同一方向更新参数, 那训练就很顺利. 不过这是个小概率事件

为了使两个输出向同一方向更新参数, 可以加载单个输出时训练好的参数来初始化网络参数. 为更新参数指明一个大概方向. 当成 “药引子”, 有了它, 训练就比较容易了. 那要去哪里找预训练模型呢, 准确地讲是训练好的参数. 这个得来全不费功夫. 保姆级 Keras 实现 Faster R-CNN 五 中保存的参数就是极好的. 如果你没有保存的话, 那就可能要重新训练了再保存了. 你要是嫌再训练一次比较麻烦. 当然也可以 下载我训练好的

假设现在你已经有 “药引子” 了, 那要怎么使用呢? 先还不急

二. 冻结主干网络

RPN 网络分类的部分相对比较容易训练, 训练好之后为了不让加入的回归部分破坏已有的参数状态, 需要将主干网络参数冻结. 因为训练分类时主干网络已经具备了特征提取的功能, 可以暂时不用更新

Keras 网络每一层都有一个 trainable 参数, 这个参数用来控制在训练的时候是否需要更新参数. 利用这个参数就可以冻结住主干网络. 我们先定义一个函数来实现这个功能, 后面还会用到

# 冻结 layers
# layers: 需要冻结的层的 name 或者序号, 由 by_name 来决定
# freeze: True 表示冻结, False 表示解冻
def freeze_layers(model, layers, freeze, by_name = False):
    if by_name:
        for name in layers:
            model.get_layer(name).trainable = not freeze
    else:
        for i in layers:
            model.get_layer(index = i).trainable = not freeze

有了这个函数, 就可以按下面的代码把主干网络冻结. 以编译模型之后加入下面代码

# 编译模型
rpn_model.compile(optimizer = keras.optimizers.Adam(learning_rate = 0.0001),
                  loss = [rpn_cls_loss, rpn_reg_loss],
                  loss_weights = [1.0, 10.0],
                  metrics = {
      "rpn_cls": rpn_cls_acc, "rpn_reg": rpn_reg_acc})

# 冻结指定层
freeze_layers(rpn_model, [i for i in range(18)], True)

for layer in rpn_model.layers:
    print("Layer name:", layer.name, "Trainable:", layer.trainable)

Layer name: input Trainable: False
Layer name: vgg16_x1_1 Trainable: False
Layer name: vgg16_x1_2 Trainable: False
Layer name: max_pooling2d Trainable: False
Layer name: vgg16_x2_1 Trainable: False
Layer name: vgg16_x2_2 Trainable: False
Layer name: max_pooling2d_1 Trainable: False
Layer name: vgg16_x3_1 Trainable: False
Layer name: vgg16_x3_2 Trainable: False
Layer name: vgg16_x3_3 Trainable: False
Layer name: max_pooling2d_2 Trainable: False
Layer name: vgg16_x4_1 Trainable: False
Layer name: vgg16_x4_2 Trainable: False
Layer name: vgg16_x4_3 Trainable: False
Layer name: max_pooling2d_3 Trainable: False
Layer name: vgg16_x5_1 Trainable: False
Layer name: vgg16_x5_2 Trainable: False
Layer name: vgg16_x5_3 Trainable: False
Layer name: rpn_conv Trainable: True
Layer name: rpn_cls Trainable: True
Layer name: rpn_reg Trainable: True

可以看到, 这时只有 rpn_conv, rpn_cls, rpn_reg 三层的参数可以被更新

三. 训练

现在就可以加载训练好的参数进行训练了

# 训练模型
project_name = "faster_rcnn"

epochs = 64
batch_size = 8
augment = AUGMENT_NUM # 一张图像增强后的图像数量, 这里是 4

train_reader = input_reader(train_set, CATEGORIES, batch_size = batch_size,
                            augment_fun = data_augment)
valid_reader = input_reader(valid_set, CATEGORIES, batch_size = batch_size,
                            augment_fun = data_augment)

# 加载训练好的参数
rpn_model.load_weights(osp.join(log_path, project_name + "_weights.h5"), True)

history = rpn_model.fit(train_reader,
                        steps_per_epoch = len(train_set) * augment // batch_size,
                        epochs = epochs,
                        verbose = 1,
                        validation_data = valid_reader,
                        validation_steps = max(1, len(valid_set) * augment // batch_size),
                        max_queue_size = 8,
                        workers = 1)

Train for 2004 steps, validate for 250 steps
Epoch 1/64
   3/2004 [..............................] - ETA: 19:32 - loss: 0.7627 - rpn_cls_loss: 0.0463 - 
   rpn_reg_loss: 0.0716 - rpn_cls_rpn_cls_acc: 0.9937 - rpn_reg_rpn_reg_acc: -0.2749

你会发现, 这样训练 Loss 下降会比较快, 也不容易卡住

至此, RPN 网络的训练部分完成

上一篇: 保姆级 Keras 实现 Faster R-CNN 七
下一篇: 保姆级 Keras 实现 Faster R-CNN 九