YOLOv1 模型构建与训练

相关文章

• YOLOv1 论文简要
• YOLOv1 数据集加载
• YOLOv1 损失函数
• YOLOv1 模型构建与训练
• YOLOv1 目标检测

项目地址：YOLOv1 VOC 2007

笔者训练的权重地址：阿里云盘分享

10 秒文章速览

本文主要讲解了 YOLOv1 的模型构建、模型训练

模型构建

对于模型的构建，我们不采用论文中的方案，而是使用 ResNet 模型。至于为什么，在笔者的观测下，ResNet练的训练速度明显更快

YOLOv1 模型

但在这里笔者还是贴出论文中的模型，向前辈致敬🫡

# 根据原论文构建的模型
def get_YOLOv1():
    model = keras.Sequential([
        keras.layers.Conv2D(64, (7, 7), 2, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.MaxPool2D((2, 2), 2),

        keras.layers.Conv2D(192, (3, 3), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.MaxPool2D((2, 2), 2),

        keras.layers.Conv2D(128, (1, 1), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(256, (3, 3), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(256, (1, 1), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(512, (3, 3), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.MaxPool2D((2, 2), 2),

        keras.layers.Conv2D(256, (1, 1), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(512, (3, 3), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(256, (1, 1), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(512, (3, 3), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(256, (1, 1), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(512, (3, 3), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(256, (1, 1), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(512, (3, 3), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(512, (1, 1), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(1024, (3, 3), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.MaxPool2D((2, 2), 2),

        keras.layers.Conv2D(512, (1, 1), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(1024, (3, 3), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(512, (1, 1), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(1024, (3, 3), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(1024, (3, 3), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(1024, (3, 3), 2, 'same'),
        keras.layers.LeakyReLU(0.1),

        keras.layers.Conv2D(1024, (3, 3), 1, 'same'),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Conv2D(1024, (3, 3), 1, 'same'),
        keras.layers.LeakyReLU(0.1),

        keras.layers.Flatten(),
        keras.layers.Dense(4096),
        keras.layers.LeakyReLU(0.1),
        keras.layers.Dropout(rate=0.5),
        keras.layers.Dense(1470),
        keras.layers.Reshape((7, 7, 30)),
        keras.layers.Activation('sigmoid')
    ])
    model.build(input_shape=(None, 224,224, 3))

    return model

ResNet50V2 模型

使用 ResNet 作为模型，还有一个原因是 TensorFlow 提供的模型可以加载 ImageNet 数据集的训练权重

在全连接部分，我们参考论文中的构建方式，只不过将 Flatten 层改为了 GlobalAvgPool2D 层

# 以 ResNet50V2 作为模型主干
ResNet = keras.applications.ResNet50V2(input_shape=(448, 448, 3), include_top=False)

# 构建全连接层部分
x = ResNet.output
x = keras.layers.GlobalAvgPool2D()(x)
x = keras.layers.Dense(4096)(x)
x = keras.layers.LeakyReLU(0.1)(x)
x = keras.layers.Dropout(rate=0.5)(x)
x = keras.layers.Dense(1470)(x)
x = keras.layers.Reshape((7, 7, 30))(x)
x = keras.layers.Activation('sigmoid')(x)

model = keras.Model(ResNet.input, x)

模型训练

对于模型的训练，同样参考了论文中的方案（有意思的是笔者在训练中也使用了 TensorFlow 提供的分段衰退的方法，但效果似乎不如下面这种简单粗暴的方法）

optimizer = keras.optimizers.SGD(learning_rate=0.0055, momentum=0.9, weight_decay=0.0005)
model.compile(optimizer=optimizer, loss=get_loss)

model.fit(train_set, epochs=1)
model.optimizer.learning_rate.assign(0.01)
model.fit(train_set, epochs=75)
model.optimizer.learning_rate.assign(0.001)
model.fit(train_set, epochs=30)
model.optimizer.learning_rate.assign(0.0001)
model.fit(train_set, epochs=30)

如果在不加载预训练权重的情况下，这里也给出笔者的训练方案

使用 Adam 优化器，weight_decay 设置为0.0005，训练步骤如下

• 以0.0003为学习率，训练75个周期
• 以0.0001为学习率，训练50个周期
• 以0.00005为学习率，训练25个周期
• 以0.00003为学习率，训练25个周期