保姆级 Keras 实现 Faster R-CNN 五

上一篇文章 我们实现了数据的增强和数据读入 Generator, 那现在能不能进行训练了呢? 还是不行. 因为还有一个特别重要也是很多人最棘手问题: 损失函数

一. 损失函数

Faster R-CNN 的损失函数跟其他常见的分类损失函数不一样. 我们入门时学的分类网络的损失函数时网络的每一个输出都是参与计算的, 而 Faster R-CNN 有一些样本标记为不参与训练, 也就是不参与损失函数计算. 那要怎样才能只计算正样本和负样本呢? 看码说话, 其实灰常简单

# RPN 网络分类损失函数
def rpn_cls_loss(y_true, y_pred):
    # mask 用于判断哪些 anchor 是参训计算, 哪些不计算
    mask = tf.cast(y_true > NEUTRAL, dtype = tf.float32)
    # 因为用的是 sigmoid 激活, 所以是 binary_crossentropy
    loss = K.sum(mask * K.binary_crossentropy(y_true, y_pred)) / (1e-6 + K.sum(mask))
    
    return loss

代码中 mask 用于标记哪些样本参与计算. 在打标签的时候, -1 为中立样本, y_true > NEUTRAL 就是把 > -1 也就是 0 或 1 的标签 位置 转换成 True, 等于 -1 的标签 位置 转换成 False. 再经过 tf.cast 函数把 True 转换成 1.0, False 转换成 0.0, 转换过后, mask 就相当于一个只有 0.0 和 1.0 的一个 9 通道张量

K.binary_crossentropy(y_true, y_pred) 这句计算的时候全部的样本(包括中立样本)都是参与计算了的, 计算完成之后也是一个 9 通道的张量. 当 mask 乘上 binary_crossentropy 的结果后, 因为是相同位置的元素相乘, 所以中立样本的位置变成了 0, K.sum 的功能就只是计算正负样本的和了. 我们知道, 计算 Loss 还要取平均, 所以 K.sum(mask) 就是参与计算的元素的个数, 这样一除 Loss 结果就出来了. 是不是相见恨晚!

二. 精度计算

同样, 因为有的样本不参与计算 Loss, 所以计算训练精度的时候就不能带上它们玩, 也要弄一个自定义的精度自己玩. 代码也是相当的简单

# 分类精度评价函数
def rpn_cls_acc(y_true, y_pred):
    # idxs 表示参与计算元素的索引
    idxs = tf.where(K.not_equal(y_true, NEUTRAL))
    
    y_true = tf.gather_nd(y_true, idxs)
    y_pred = tf.gather_nd(y_pred, idxs)
    
    acc = keras.metrics.binary_accuracy(y_true, y_pred, threshold = 0.5)
    return acc

三. 模型编译

模型编译的时候要指定所用的损失函数和精度评估函数

# 编译模型
project_name = "faster_rcnn"

rpn_model.compile(optimizer = keras.optimizers.Adam(learning_rate = 1e-4),
                  loss = rpn_cls_loss,
                  metrics = [rpn_cls_acc])

Keras 的 Adam 默认 learning_rate 是 0.001, 经过我的试验, 选 0.0001 比较好训练一点

四. 模型训练

终于到了你期待的训练时刻了

# 训练模型
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)

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)

训练的代码按套路来的, 所以也没有什么好讲的. 如果你的显存比较小, 就要把 batch_size 设置小一点, 不然会爆炸(OOM)

Train for 2004 steps, validate for 250 steps
Epoch 1/64
  74/2004 [=>............................] - ETA: 22:51 - loss: 3.4566 - rpn_cls_acc: 0.9460

到这里, 你就等待训练完成吧… 看样子没有一天一夜是训练不完了. 训练完成之后保存模型和参数

# 保存模型与参数
rpn_model.save(osp.join(log_path, project_name + "_model.h5"))
rpn_model.save_weights(osp.join(log_path, project_name + "_weights.h5"))

五. 代码下载

示例代码可下载 Jupyter Notebook 示例代码

上一篇: 保姆级 Keras 实现 Faster R-CNN 四
下一篇: 保姆级 Keras 实现 Faster R-CNN 六