【YOLO系列】YOLOv3代码详解(三)：训练脚本train.py

前言

????????以下内容仅为个人在学习人工智能中所记录的笔记，先将目标识别算法yolo系列的整理出来分享给大家，供大家学习参考。

????????本文仅对YOLOV3代码中关键部分进行了注释，未掌握Python基础的铁汁可以自己百度一下。

????????若文中内容有误，希望大家批评指正。

---

资料下载

????????YOLOV3论文下载地址：YOLOv3：An Incremental Improvement

回顾

????????YOLO V1：【YOLO系列】YOLO V1论文思想详解

????????YOLO V2：【YOLO系列】YOLO V2论文思想详解

????????YOLO V3：【YOLO系列】 YOLOv3论文思想详解

项目地址

????????YOLO V3 keras版本：下载地址

????????YOLO V3 Tensorflow版本：下载地址

????????YOLO V3 Pytorch版本：下载地址

Gitee仓库

????????YOLO V3 各版本：yolov3各版本

---

YOLO V3代码详解

????????YOLO V3代码详解（一）：【YOLO系列】YOLOv3代码详解(一)：主脚本yolo_video.py

????????YOLO V3代码详解（二）：【YOLO系列】YOLOv3代码详解(二)：检测脚本yolo.py

---

????????本文主要基于keras版本进行讲解

????????话不多说，直接上代码

---

一、代码详解

1、_main()函数

????????（1）设置训练集文件、日志文件地址、数据集标签文件、Anchor Box文件，加载类别名称、类别数量、Anchor Box以及输入图片的默认尺寸；

????????（2）根据Anchor Box的数量判断创建什么model

????????（3）创建TensorBoard回调，存储日志文件；创建模型检测点ModelCheckpoint，检查模型的性能，判断是否保存权重；创建自动调整LR平台ReduceLROnPlateau，防止模型在训练过程中陷入局部最小值；创建自动停止训练函数EarlyStopping，验证损失不再改善时提前停止训练模型；

????????（4）划分训练集与验证集；

????????（5）编译模型，训练并保存模型。

def _main():
    # 加载训练集文件
    annotation_path = 'train.txt'
    # 加载日志文件地址
    log_dir = 'logs/000/'
    # 加载数据集标签类别文件
    classes_path = 'model_data/voc_classes.txt'
    # 加载聚类生成的Anchor Box文件
    anchors_path = 'model_data/yolo_anchors.txt'
    # 获取各标签类别名称
    class_names = get_classes(classes_path)
    # 获取标签类别数量
    num_classes = len(class_names)
    # 获取Anchor Box
    anchors = get_anchors(anchors_path)

    # 默认输入图片尺寸为[416, 416]
    input_shape = (416, 416)  # multiple of 32, hw

    # 如果Anchor Box数量为6个，则为tiny版本的Yolo v3，否则为正常版本
    is_tiny_version = len(anchors) == 6  # default setting
    if is_tiny_version:
        model = create_tiny_model(input_shape, anchors, num_classes,
                                  freeze_body=2, weights_path='model_data/tiny_yolo_weights.h5')
    else:
        model = create_model(input_shape, anchors, num_classes,
                             freeze_body=2, weights_path='model_data/yolo_weights.h5')  # make sure you know what you freeze

    # 创建一个TensorBoard回调，存储日志文件
    # TensorBoard是一个非常有用的回调（callback），它允许你在训练过程中和训练结束后监控模型的性能。
    # 通过TensorBoard，你可以可视化诸如损失、准确率、权重和激活等指标。
    logging = TensorBoard(log_dir=log_dir)
    # 回调在每个 epoch 结束后检查模型的性能，如果模型在验证集上的性能有所提高，那么它将保存模型的权重。
    # log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5' 用于保存最佳模型的权重
    # monitor='val_loss' 参数指定了想要监控的指标（验证损失），
    # save_best_only=True 表示只有当监控的指标有所改善时，才保存模型的权重。
    # period=3，那么 ModelCheckpoint 将在每个 3 个 epochs 结束后保存模型的权重。
    checkpoint = ModelCheckpoint(log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5',
                                 monitor='val_loss', save_weights_only=True, save_best_only=True, period=3)
    # ReduceLROnPlateau用于在验证损失不再改善时减少学习率。这对于防止模型在训练过程中陷入局部最小值并帮助模型更好地收敛非常有用。
    # monitor: 你想要监视的指标。('val_loss')。
    # factor: 当监视的指标不再改善时，用于减少学习率的比例。这里学习率将乘以 0.1。
    # patience: 等待监视的指标不再改善的 epoch 数。这里等待 3 个 epochs。
    # verbose: 设置为 1 时，回调将在每个 epoch 后输出有关学习率的信息。
    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=3, verbose=1)
    # EarlyStopping 用于在验证损失不再改善时提前停止训练模型。这有助于防止模型过拟合，并节省计算资源。
    # monitor: 你想要监视的指标。在这个例子中，我们监视验证损失 ('val_loss')。
    # min_delta: 用于确定指标是否有所改善的阈值。如果两个连续的 epochs 中的损失之间的差异小于这个值，那么它将认为损失已经不再改善。
    # patience: 等待监视的指标不再改善的 epoch 数。在这里将等待 10 个 epochs。
    # verbose: 设置为 1 时，回调将在提前停止训练时输出一条消息。
    early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=10, verbose=1)

    # 验证集比例
    val_split = 0.1
    # 加载train集
    with open(annotation_path) as f:
        lines = f.readlines()
    np.random.seed(10101)
    np.random.shuffle(lines)
    np.random.seed(None)
    # 设置验证集的数量
    num_val = int(len(lines) * val_split)
    # 设置train集的数据
    num_train = len(lines) - num_val

    # Train with frozen layers first, to get a stable loss.
    # Adjust num epochs to your dataset. This step is enough to obtain a not bad model.
    if True:
        # 编译模型,使用Adam优化器,学习率为0.001
        model.compile(optimizer=Adam(lr=1e-3), loss={
            # use custom yolo_loss Lambda layer.
            'yolo_loss': lambda y_true, y_pred: y_pred})

        # 分为32批次
        batch_size = 32
        print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))
        # 传递数据，训练模型
        # data_generator_wrapper()是一个数据生成器，返回[image_data, *y_true], np.zeros(batch_size)
        # image_data：通过数据增强方式得到的图片数据
        model.fit_generator(data_generator_wrapper(lines[:num_train], batch_size, input_shape, anchors, num_classes),
                            steps_per_epoch=max(1, num_train // batch_size),
                            validation_data=data_generator_wrapper(lines[num_train:], batch_size, input_shape, anchors, num_classes),
                            validation_steps=max(1, num_val // batch_size),
                            epochs=50,
                            initial_epoch=0,
                            callbacks=[logging, checkpoint])
        # 保存模型
        model.save_weights(log_dir + 'trained_weights_stage_1.h5')

    # Unfreeze and continue training, to fine-tune.
    # Train longer if the result is not good.
    # 训练模型的所有层参数
    if True:
        for i in range(len(model.layers)):
            model.layers[i].trainable = True
        model.compile(optimizer=Adam(lr=1e-4), loss={'yolo_loss': lambda y_true, y_pred: y_pred})  # recompile to apply the change
        print('Unfreeze all of the layers.')

        batch_size = 32  # note that more GPU memory is required after unfreezing the body
        print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))
        model.fit_generator(data_generator_wrapper(lines[:num_train], batch_size, input_shape, anchors, num_classes),
                            steps_per_epoch=max(1, num_train // batch_size),
                            validation_data=data_generator_wrapper(lines[num_train:], batch_size, input_shape, anchors, num_classes),
                            validation_steps=max(1, num_val // batch_size),
                            epochs=100,
                            initial_epoch=50,
                            callbacks=[logging, checkpoint, reduce_lr, early_stopping])
        model.save_weights(log_dir + 'trained_weights_final.h5')

    # Further training if needed.

2、获取标签类别名称与Anchor Box

# 获取class
def get_classes(classes_path):
    '''loads the classes'''
    with open(classes_path) as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]
    return class_names


# 获取anchor box
def get_anchors(anchors_path):
    '''loads the anchors from a file'''
    with open(anchors_path) as f:
        anchors = f.readline()
    anchors = [float(x) for x in anchors.split(',')]
    # 返回N行2列的二维数组
    return np.array(anchors).reshape(-1, 2)

3、定义创建model函数

????????（1）创建计算图、定义输入图片的Tensor、y_true的形状

????????（2）判断是否加载预训练模型的权重yolo_weights.h5，并冻结模型中除3层输出层以外的所有层；

????????（3）定义模型的loss

????????（4）返回model

def create_model(input_shape, anchors, num_classes, load_pretrained=True, freeze_body=2,
                 weights_path='model_data/yolo_weights.h5'):
    '''create the training model'''
    # 创建模型
    # 首先清楚其他的计算图
    K.clear_session()  # get a new session
    # 定义一个输入层，接受任意大小和形状的彩色（3通道）图片，返回一个Tensor
    image_input = Input(shape=(None, None, 3))
    # 获取输入图片的高和宽
    h, w = input_shape
    # 获取Anchor Box的数量
    num_anchors = len(anchors)

    # 定义三个输出层数组y_true，对应输出的3个特征图(13,13),(26,26),(52,52),即输入的图片高和宽被32、16、8整除
    # 同时定义了Anchor Box的数量，num_classes+5表示标签的类别数量+5,5表示4（bbox的offset）+1（是否为物体）
    # 在YOLOv3中，最后会在不同尺度的特征图上得到了一个N×N×[3×(4+1+80)]的tensor，就与之对应
    y_true = [Input(shape=(h // {0: 32, 1: 16, 2: 8}[l], w // {0: 32, 1: 16, 2: 8}[l], num_anchors // 3, num_classes + 5)) for l in range(3)]

    # 创建YoloV3 模型
    model_body = yolo_body(image_input, num_anchors // 3, num_classes)
    print('Create YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes))

    # 是否加载预训练yoloV3模型的权重yolo_weights.h5，并冻结模型中除3层输出层以外的所有层
    # 冻结模型中除输出层以外的所有层，意味着在训练过程中，这些层不会进行参数更新。主要为减少计算复杂度，提高泛化能力
    # 需要注意的是，冻结模型中除输出层以外的所有层并不意味着这些层的参数就完全不变。在某些情况下，可能需要对这些层进行微调（fine-tuning），以适应新的目标检测任务。
    if load_pretrained:
        model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)
        print('Load weights {}.'.format(weights_path))
        if freeze_body in [1, 2]:
            # Freeze darknet53 body or freeze all but 3 output layers.
            # 185 表示在yoloV3中总共有185层
            # 在此是冻结模型中除3层输出层以外的所有层，这些层trainable=False，不训练
            num = (185, len(model_body.layers) - 3)[freeze_body - 1]
            for i in range(num):
                model_body.layers[i].trainable = False
            print('Freeze the first {} layers of total {} layers.'.format(num, len(model_body.layers)))

    # 定义模型的loss，使用Keras.layers.Lambda()函数，传递一个name为yolo_loss的损失函数，输入数组形状为(1,)
    # 这里的model设置的阈值为0.5
    model_loss = Lambda(yolo_loss, output_shape=(1,), name='yolo_loss',
                        arguments={'anchors': anchors, 'num_classes': num_classes, 'ignore_thresh': 0.5})(
        [*model_body.output, *y_true])
    model = Model([model_body.input, *y_true], model_loss)

    return model


def create_tiny_model(input_shape, anchors, num_classes, load_pretrained=True, freeze_body=2,
                      weights_path='model_data/tiny_yolo_weights.h5'):
    '''为tiny YoloV3创建训练模型'''
    K.clear_session()  # 创建一个新的计算图，即删除当前会话中的所有变量和模型，并释放与该会话关联的资源。
    # 定义一个输入层，接受任意大小和形状的彩色（3通道）图片，返回一个Tensor
    image_input = Input(shape=(None, None, 3))
    # 获取输入图片的高和宽
    h, w = input_shape
    # 获取Anchor Box的数量
    num_anchors = len(anchors)

    # 定义两个输出层y_true，对应输出的3个特征图(13,13),(26,26),即输入的图片高和宽被32、16整除
    # 同时定义了Anchor Box的数量，num_classes+5表示标签的类别数量+5,5表示4（bbox的offset）+1（是否为物体）
    # 在YOLOv3中，最后会在不同尺度的特征图上得到了一个N×N×[3×(4+1+80)]的tensor，就与之对应
    y_true = [Input(shape=(h // {0: 32, 1: 16}[l], w // {0: 32, 1: 16}[l], num_anchors // 2, num_classes + 5)) for l in range(2)]

    # 创建tiny YoloV3 模型
    model_body = tiny_yolo_body(image_input, num_anchors // 2, num_classes)
    print('Create Tiny YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes))

    # 是否加载预tiny yoloV3模型的权重tiny_yolo_weights.h5，并冻结模型中除2层输出层以外的所有层
    # 冻结模型中除输出层以外的所有层，意味着在训练过程中，这些层不会进行参数更新。主要为减少计算复杂度，提高泛化能力
    # 需要注意的是，冻结模型中除输出层以外的所有层并不意味着这些层的参数就完全不变。在某些情况下，可能需要对这些层进行微调（fine-tuning），以适应新的目标检测任务。
    if load_pretrained:
        model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)
        print('Load weights {}.'.format(weights_path))
        if freeze_body in [1, 2]:
            # Freeze the darknet body or freeze all but 2 output layers.
            # 20 表示在tiny yoloV3中总共有20层
            # 在此是冻结模型中除2层输出层以外的所有层，这些层trainable=False
            num = (20, len(model_body.layers) - 2)[freeze_body - 1]
            for i in range(num): model_body.layers[i].trainable = False
            print('Freeze the first {} layers of total {} layers.'.format(num, len(model_body.layers)))

    # 定义模型的loss，使用Keras.layers.Lambda()函数，传递一个name为yolo_loss的损失函数，输入数组形状为(1,)
    # 这里的tiny model设置的阈值为0.7
    model_loss = Lambda(yolo_loss, output_shape=(1,), name='yolo_loss',
                        arguments={'anchors': anchors, 'num_classes': num_classes, 'ignore_thresh': 0.7})(
        [*model_body.output, *y_true])
    model = Model([model_body.input, *y_true], model_loss)

    return model

4、数据生成器函数

????????（1）调用get_random_data()函数实现图片数据增强；

????????（2）预先计算Anchor box与GT框的IOU，将最大的IOU对应的GT框的y_true记为1；

????????（3）返回一个数据生成器

def data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes):
    '''data generator for fit_generator'''
    n = len(annotation_lines)
    i = 0
    while True:
        image_data = []
        box_data = []
        for b in range(batch_size):
            if i == 0:
                # 开始前先打乱train集
                np.random.shuffle(annotation_lines)
            # 调用get_random_data()函数实现图片数据增强
            image, box = get_random_data(annotation_lines[i], input_shape, random=True)
            image_data.append(image)
            box_data.append(box)
            i = (i + 1) % n
        image_data = np.array(image_data)
        box_data = np.array(box_data)
        # 预先计算Anchor box与GT框的IOU，将最大的IOU对应的GT框的y_true记为1
        y_true = preprocess_true_boxes(box_data, input_shape, anchors, num_classes)
        yield [image_data, *y_true], np.zeros(batch_size)


def data_generator_wrapper(annotation_lines, batch_size, input_shape, anchors, num_classes):
    n = len(annotation_lines)
    if n == 0 or batch_size <= 0: return None
    return data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes)