☰
用TensorFlow手写YOLOv1目标检测:从标签编码到损失函数与推理避坑
2026/10/12 21:42:04 网站建设 项目流程

简介:YOLOv1作为实时目标检测的开山之作,将检测任务转化为回归问题,通过单次前向推理直接输出边界框与类别概率。该压缩包提供带注释的完整TensorFlow实现,共28个文件,包含9个Python源文件、12个pyc编译缓存文件,以及测试图片、授权说明和README文档,整体仅302KB,轻量易部署。源码覆盖模型定义、VOC数据加载、训练与测试脚本、配置参数等核心模块,并附数据下载与计时辅助工具,结构清晰,适合入门者逐模块研读。已有992人学习下载,对希望动手复现YOLOv1模型或熟悉TensorFlow目标检测流程的开发者,是一份可直接运行的参考实现。

1. 用 TensorFlow 手写 YOLOv1,先想清楚在复现什么

用 TensorFlow 复现 YOLOv1,本质上不是「用深度学习框架实现一个检测算法」,而是把「目标检测被重新定义成回归问题」这件事亲手做一遍:输入一张 448×448 的图,输出一个 7×7×30 的张量,所有框、类别、置信度全部从卷积网络里直接回归出来。这个思路今天看很朴素,但它的损失函数设计和标签编码方式,影响了后面整个单阶段检测家族。

这个标题背后的人群很明确:不满足于调接口、想真正看懂检测原理的人,或者正在做课程实验、面试项目、想转入检测方向的一线开发。复现的价值不在于把 mAP 刷到多高,而在于你把数据编码、网络结构、损失函数、NMS 这四块全部亲手串起来一次,以后再去看 SSD、YOLOv2、RetinaNet,理解成本会直线下降。你会很快发现,剩下的百分之六七十的精力其实不在模型本身,而在数据管线和损失函数的细节里。

2. 数据准备:从标注框到 7×7×30 的标签张量

2.1 标签编码:YOLOv1 的 ground truth 长什么样

YOLOv1 把输入图划分成 7×7 的网格,每个网格负责预测 2 个边界框和一组类别概率。最终输出的通道数是 30,拆开看就是 2×5+20:每个框有中心坐标 x、y,宽高 w、h,和一个置信度,然后整格共享 20 个类别的条件概率。先记住这个编码,后面所有代码都在围绕它转。

标签编码的第一步是归一化。所有坐标都以整张图宽高为基准归一化到 0~1,中心点还要进一步转成「相对于所在网格」的坐标。宽高保持相对整图的归一化值,这样在 0~1 之间,更重要的是训练时避免大图小图比例不一致导致回归不稳定。这一步一定要归一化,不然损失函数里的坐标项会被像素量大数值带偏,整个训练直接失控。

第二个关键点是「谁负责这个目标」。每个物体只会落到一个网格里,由中心点所在网格负责;该网格的两个预测框里,训练时会计算两个框和真实框的 IoU,IoU 更大的那个被标记为 responsible box,只有它参与坐标回归和置信度回归。另一个框如果 IoU 低于阈值,就进入 noobj 分支。

把 VOC 格式的标注从 xml 解析出来,转换成 7×7×30 张量的过程,我习惯用一个 Python 函数一次做完,输出固定长度的 float 数组,供后面写 TFRecord 用。

# 把单个标注框转换进 7*7*30 的 label 张量 import numpy as np NUM_GRID = 7 NUM_BOX = 2 NUM_CLASS = 20 # VOC 格式数据集的默认类别数 def encode_label(boxes, classes, img_w, img_h): # boxes: [(x1, y1, x2, y2), ...],像素坐标 # classes: [int, ...],与 boxes 一一对应 label = np.zeros((NUM_GRID, NUM_GRID, 5 * NUM_BOX + NUM_CLASS), dtype=np.float32) for box, cls in zip(boxes, classes): x1, y1, x2, y2 = box cx = (x1 + x2) / 2.0 cy = (y1 + y2) / 2.0 w = x2 - x1 h = y2 - y1 # 归一化到图片尺寸 cx, cy, w, h = cx / img_w, cy / img_h, w / img_w, h / img_h # 中心点落到哪个网格 grid_x = int(cx * NUM_GRID) grid_y = int(cy * NUM_GRID) grid_x = min(grid_x, NUM_GRID - 1) grid_y = min(grid_y, NUM_GRID - 1) # 相对本网格的坐标,范围约 0~1 x_cell = cx * NUM_GRID - grid_x y_cell = cy * NUM_GRID - grid_y # 第一个框的 5 个值,class 概率放在 20 维里 cell = label[grid_y, grid_x] if cell[0] == 0: cell[0:5] = [x_cell, y_cell, w, h, 1.0] else: # 同一网格出现第二个目标时,放进第二个框,若也冲突则覆盖概率不处理 if cell[5] == 0: cell[5:10] = [x_cell, y_cell, w, h, 1.0] cell[10 + cls] = 1.0 return label

这段代码把每个目标的 x1y1x2y2 转成归一化的中心宽高,再放进对应网格。cell[0] == 0用来判断第一个框是否被占用,是因为一个网格原则上只负责一个目标,多目标挤在一个网格里时直接放到第二个框位。

注意 x_cell、y_cell 用的是「相对网格的浮点坐标」,不是绝对坐标。训练时网络输出的 xy 也要这样对齐,否则损失函数在空间上对不齐。类别概率在cell[10 + cls]处置 1,同一网格的多目标共享同一份类别向量,这也是 YOLOv1 对小目标拥挤场景效果差的直接原因——一个格子只有一组类别概率,物理上就没法区分叠在一起的多个类。

2.2 把图片和标签写进 TFRecord:避免训练时内存爆炸

目标检测训练通常要跑几十轮,图片总量一大,全部放内存会非常吃力。常见做法是先把处理好的图片和标签张量序列化进 TFRecord,训练时用tf.data做流式读取。这张图解码、缩放、归一化的开销全部放在数据管线里,GPU 利用率会明显改善。

import tensorflow as tf def write_tfrecord(image_paths, labels, record_path): with tf.io.TFRecordWriter(record_path) as writer: for img_path, label in zip(image_paths, labels): img_bytes = tf.io.read_file(img_path).numpy() label_bytes = label.astype(np.float32).tobytes() feature = { 'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_bytes])), 'label': tf.train.Feature(bytes_list=tf.train.BytesList(value=[label_bytes])), } example = tf.train.Example(features=tf.train.Features(feature=feature)) writer.write(example.SerializeToString())

标签提前编码成 7×7×30 的稠密数组,写进 record。图片保留 JPEG 原始字节,由读取端解码。这样磁盘占用小,读取时还可以直接基于原始编码做随机裁剪或翻转,不必先解码成 numpy 数组再折腾。

读取端的解析函数要注意两点:第一,tf.io.decode_raw需要明确指定 dtype 为tf.float32,不然默认解析成 uint8,张量的数值全乱了;第二,解析完必须tf.reshape(label, (7, 7, 30)),因为序列化的时候它已经是一维字节流,reshape 的顺序是按行优先恢复的,正好对应原来的网格坐标。

def parse_example(example_proto): feature_desc = { 'image': tf.io.FixedLenFeature([], tf.string), 'label': tf.io.FixedLenFeature([], tf.string), } parsed = tf.io.parse_single_example(example_proto, feature_desc) image = tf.image.decode_jpeg(parsed['image'], channels=3) label = tf.io.decode_raw(parsed['label'], tf.float32) label = tf.reshape(label, (NUM_GRID, NUM_GRID, 5 * NUM_BOX + NUM_CLASS)) return image, label

解析出来的图片尺寸是不固定的,后面还要统一 resize 到 448×448。常见做法是在dataset.map里顺手做掉,不要等进模型再处理,否则动态尺寸会拖慢图模式执行效率。

2.3 数据增强:翻转时最容易把 label 改坏

YOLOv1 训练里最简单有效的增强是水平翻转和随机亮度饱和度抖动。水平翻转在tf.image.flip_left_right之后,标注里的中心点 x 坐标要变成1 - x。grid_x 也会镜像到对面网格,x_cell 同样要变成1 - x_cell。很多人在这一步翻车,翻转了图却忘了改 label,训练 loss 居高不下还以为是模型问题。

def augment(image, label, is_train): image = tf.image.resize(image, (448, 448)) image = tf.cast(image, tf.float32) / 255.0 if is_train: if tf.random.uniform(()) > 0.5: image = tf.image.flip_left_right(image) label = flip_label_horizontal(label) image = tf.image.random_brightness(image, max_delta=0.05) image = tf.image.random_saturation(image, lower=0.8, upper=1.2) return image, label

flip_label_horizontal里面单独处理 7×7 上的每个网格,把 xy 坐标镜像,然后交换网格列索引,这里最容易写错的顺序是先改坐标再换格子。如果先交换列再改 x_cell,结果会同时错两处,且叠加起来看起来像偶发错误,排查很久才能发现。

另外,彩色抖动在灰度图上会报错或毫无意义,训练前检查一下数据集里有没有灰度图。直接tf.image.random_saturation会强行把灰度图当作 3 通道处理,虽然不会崩,但颜色增强起不到预期效果。

3. 网络结构:用 tf.keras 把 YOLOv1 搭起来

3.1 输出张量结构解读:为什么是 7×7×30 而不是别的

网络结构本身不算深:24 层卷积加 2 层全连接,最后输出 1470 个数,reshape 成 7×7×30。30 这个数字前面已经聊过:2 个框,每个 5 个参数,加上 20 类概率。7×7 的选择是原版在 VOC 格式数据集上定下的,复现时没必要改,改了还要重写整个 ground truth 生成逻辑。

这个结构的精髓在于「检测头就是回归头」。网络没有区域提议、没有锚框、没有后处理网络分支,所有预测都直接从最后的特征图出来。理解这一点,才能明白为什么后续 YOLO 系列改进都在「让每个位置预测得更准、更多」这条线上做文章,而不是重新设计两阶段流程。

主干部分的卷积设计有很多1×1卷积穿插在3×3卷积之间,目的是降通道数、减少计算量,这也是后来瓶颈结构的前身。全连接层在这个任务里实际承担了「全局建模」的作用,让网络能利用整图上下文判断目标位置,代价是输入尺寸必须固定。

3.2 用 Keras 函数式 API 搭建主干与检测头

import tensorflow as tf from tensorflow.keras import layers def conv_block(x, filters, kernel, stride=1): x = layers.Conv2D(filters, kernel, strides=stride, padding='same')(x) x = layers.BatchNormalization()(x) return layers.LeakyReLU(alpha=0.1)(x) def build_yolov1(num_grid=7, num_boxes=2, num_classes=20): inputs = tf.keras.Input(shape=(448, 448, 3)) x = conv_block(inputs, 64, 7, stride=2) x = layers.MaxPooling2D(2, 2)(x) x = conv_block(x, 192, 3) x = layers.MaxPooling2D(2, 2)(x) x = conv_block(x, 128, 1) x = conv_block(x, 256, 3) x = conv_block(x, 256, 1) x = conv_block(x, 512, 3) x = layers.MaxPooling2D(2, 2)(x) # 原版在这里堆了 4 组 1x1 + 3x3 卷积 for _ in range(4): x = conv_block(x, 256, 1) x = conv_block(x, 512, 3) x = conv_block(x, 512, 1) x = conv_block(x, 1024, 3) x = layers.MaxPooling2D(2, 2)(x) # 后半段保持 7x7 空间分辨率 x = conv_block(x, 512, 1) x = conv_block(x, 1024, 3) x = conv_block(x, 1024, 3, stride=2) x = conv_block(x, 1024, 3) x = layers.Flatten()(x) x = layers.Dense(4096)(x) x = layers.LeakyReLU(alpha=0.1)(x) x = layers.Dense(num_grid * num_grid * (num_boxes * 5 + num_classes))(x) outputs = layers.Reshape((num_grid, num_grid, num_boxes * 5 + num_classes))(x) return tf.keras.Model(inputs, outputs)

注意我在conv_block里加了 BatchNormalization,这一步不是原版结构,原版只在最后全连接前有 dropout,卷积层没有 BN。但直接用原版结构复现时,训练收敛很慢,对学习率极其敏感。所以常见的复现做法是加 BN 换稳定,代价是如果没有原版预训练权重,从头训练也能顺利收敛。如果你打算加载别人转好的预训练权重,网络结构必须和对方完全一致,BN 加不加、加在哪都要逐层核对。

最后一层全连接没有接任何激活函数,是线性输出。原因在于 YOLOv1 把检测值当回归问题,宽高可以大于 1(目标可以比网格大),类别概率也需要原始 logits 而不是 sigmoid 之后的值。你在实现时如果手滑给最后一层加了activation='sigmoid',训练初期看起来 loss 挺低,但 recall 会非常差,因为框的范围被锁死了。

Reshape 这里有个容易忽略的点:Flatten之后接Dense得到的是长度为 1470 的向量,再Reshape((7, 7, 30))是按行优先填充的,也就是第一个 30 属于第 (0, 0) 号网格,这和标签编码时label[grid_y, grid_x]的索引顺序一致,不会出现转置错位。

3.3 预训练权重:能加载就别从头训

从一个随机初始化开始训 YOLOv1 是血泪之路,前面几十轮 loss 高得离谱,框全缩在图像角落,看起来完全没学会。原版做法是先在分类任务上预训练前 20 层卷积,再把分类头换掉接检测头。复现时最省力的方案是用公开转换好的权重文件。

加载预训练权重最常见的坑是层名对不上。TF 的权重文件按层名索引,Darknet 转出来的权重很多层名是conv_0这类自动编号,和你 Keras 里的名字完全不一样。我一般会打印每一层的名字和权重形状,写一个手动映射字典,逐层赋值。

model = build_yolov1() # 假设 pretrained 是这个映射后的权重 dict,key 为层名 for layer in model.layers: if layer.name in pretrained: layer.set_weights(pretrained[layer.name]) print(f"loaded {layer.name}")

如果你只想快速跑通完整流程,不追求和原版完全一致的 mAP,可以直接用上面的结构从头训练,把学习率调小到 1e-4,batch size 用 16 或 32,在小数据集上跑几十轮也能看到框逐渐变准。但要做正经对比实验、要复现论文指标,预训练权重几乎是必须的。

4. 损失函数:YOLOv1 最核心的五项设计

4.1 损失拆解:坐标、宽高、置信度、类别,权重不一样

YOLOv1 的损失函数是五项之和,每一项都带着自己的权重系数。坐标误差和宽高误差权重是 5,因为坐标直接决定框的位置,必须优先学准;不含目标的置信度误差权重是 0.5,因为绝大多数网格里根本没有目标,如果不压低权重,网络会学成「永远预测没有目标」的偷懒解。

宽高误差用的是平方根差。这个设计的直觉是:同样差 10 个像素,对小框影响远比大框严重,开平方后让小框的误差在数值上被放大,梯度也会更大,网络就会更努力去修正小框。这在今天看有点粗糙,但在当时是对「大小框公平性」的一次有效探索。

置信度部分要注意 y_pred 的置信度是线性输出值,y_true 的置信度则取「预测框和真实框的 IoU」。就是说训练目标不是简单设为 1,而是这个框自己和目标的实际重合程度。这样网络学出来的置信度天然带有定位质量信息,推理时可以直接用来做 NMS 排序。

类别损失只在一个网格内有目标时计算,而且类别概率是网格共享的,和两个框无关。这就是 YOLOv1 对密集目标无能为力的结构性原因,损失函数层面没有任何一个位置能表达「同一个格子有两个不同类」。

4.2 TensorFlow 实现:mask 是怎么生成的

损失实现的核心是生成三个 mask:responsible mask、noobj mask、object mask。responsible mask 决定哪个预测框参与坐标和置信度回归,noobj mask 决定哪些框要被压到 0 置信度。

def yolo_loss(y_true, y_pred, grid=7, boxes=2, classes=20, lambda_coord=5.0, lambda_noobj=0.5): # y_true, y_pred: (batch, 7, 7, 30) pred_boxes = tf.reshape(y_pred[..., :10], (-1, grid, grid, boxes, 5)) pred_xy = pred_boxes[..., 0:2] # 相对网格的坐标 pred_wh = pred_boxes[..., 2:4] # 相对整图的宽高 pred_conf = pred_boxes[..., 4:5] true_boxes = tf.reshape(y_true[..., :10], (-1, grid, grid, boxes, 5)) true_xy = true_boxes[..., 0:2] true_wh = true_boxes[..., 2:4] true_conf = true_boxes[..., 4:5] # 转成框的绝对坐标计算 IoU pred_box_abs = xywh_to_xyxy(offset_xy_to_abs(pred_xy, pred_wh)) true_box_abs = xywh_to_xyxy(offset_xy_to_abs(true_xy, true_wh)) iou = compute_iou(pred_box_abs, true_box_abs) # (batch, 7, 7, 2) # responsible mask: 选 IoU 最大的那个框 iou_max = tf.reduce_max(iou, axis=-1, keepdims=True) resp_mask = tf.cast(tf.equal(iou, iou_max), tf.float32) # noobj mask: IoU 低于阈值且不负责的框 noobj_mask = tf.cast((iou < 0.5) & (tf.equal(resp_mask, 0.0)), tf.float32) xy_loss = lambda_coord * tf.reduce_sum(resp_mask * (pred_xy - true_xy) ** 2) wh_loss = lambda_coord * tf.reduce_sum(resp_mask * (tf.sqrt(pred_wh) - tf.sqrt(true_wh)) ** 2) conf_loss_pos = tf.reduce_sum(resp_mask * (pred_conf - iou) ** 2) conf_loss_neg = lambda_noobj * tf.reduce_sum(noobj_mask * pred_conf ** 2) class_loss = tf.reduce_sum(true_conf * (y_pred[..., 10:] - y_true[..., 10:]) ** 2) return (xy_loss + wh_loss + conf_loss_pos + conf_loss_neg + class_loss) / tf.cast(tf.shape(y_true)[0], tf.float32)

resp_mask用的是数值比较而不是tf.argmax,是为了保持和iou张量同形状,方便后面直接乘到预测框的每个维度上。如果两个框 IoU 完全相同,tf.equal会把两个框都标成 responsible,实际操作中概率极低,但如果你要精确复现原版「只选一个」,可以改成tf.one_hot(tf.argmax(iou, axis=-1), boxes)。

noobj_mask里那个&是关键:IoU 小于 0.5 且不是 responsible 的框才会被惩罚。如果少了这个与条件,responsible 框在 IoU 比较低的时候会被同时拉向两个方向——一个梯度要求它置信度接近 iou,另一个要求它置信度归零,训练会反复震荡。

类别损失里的true_conf是 (batch, 7, 7, 1),直接乘到 (batch, 7, 7, 20) 上,靠广播把没目标的网格全部清零。这里不能用reduce_mean代替reduce_sum,因为坐标项和类别项的比例已经被 lambda 系数调过,再用 mean 会把整体量纲改掉,调好的权重系数就失效了。

4.3 几个让损失函数失效的常见写法

第一,把pred_wh也取平方根然后计算(sqrt(pred_wh) - sqrt(true_wh))^2,这是对的;但有人会对 pred_wh 做tf.clip_by_value防止负数,clip 的位置必须在开平方之前,而且下界要设成 1e-6 而不是 0,因为sqrt(0)的梯度是无穷大,训练直接出 nan。

第二,xy 坐标如果用 sigmoid 激活,损失函数里的pred_xy就必须也跟着改:要先用 sigmoid 再参与计算,但 y_true 的 xy 已经是 0~1 的相对坐标,直接算平方差没问题。原版是线性输出,你加 sigmoid 属于改动,不是 bug,但改了这一处,前面预训练权重就不能完整复用了。

第三,这个损失函数返回的是整个 batch 的总和除以 batch size。如果你用tf.keras.losses.Loss子类实现,注意别在 call 里再套一层reduce_mean,二次平均会把 loss 压到一个失真的小数值,观察学习率曲线时会被误导。

5. 训练与推理避坑指南:loss 不降、框乱飞、NMS 失效

5.1 优化器与学习率:谁在让 loss 不降反升

YOLOv1 从头训练时对学习率非常敏感,经验上 Adam 用 1e-3 起步很容易发散,常见做法是先用 1e-4 让网络稳定下来,前 10 轮之后再看损失曲线决定要不要调大。SGD 加动量是另一条可靠路线,但需要配合 warmup,头几个 epoch 学习率从 1e-5 线性升到 1e-3,不然前期梯度方向太乱。

判别标准是看 loss 行为的三个阶段:第一阶段 loss 在 1e2 量级震荡,是正常的,网络在学「哪里没有目标」;第二阶段 loss 开始往下走但很慢,对应坐标项逐渐生效;第三阶段如果 loss 卡住不动,多半不是学习率问题,而是 noobj 权重和正负样本比例失衡,可以先把lambda_noobj从 0.5 降到 0.3 试试。

在tf.keras里实现 warmup 的常见写法是把学习率用tf.keras.optimizers.schedules.PolynomialDecay包一层,初始学习率initial_learning_rate=1e-5,end_learning_rate=1e-4,再配合decay_steps覆盖前 10 轮。直接手写回调控制optimizer.lr也可以,但容易在 TF2 的图执行模式下失效。

5.2 推理阶段:解码输出再 NMS,别直接用网络输出画框

训练结束后拿到的模型输出还是 7×7×30 的原始张量,直接画框会有大量重叠框。因为每个目标附近会有多个网格同时预测到它,而且同一网格的两个框可能都在往目标上靠。推理时要先做三件事:过滤低置信度、把框坐标从网格相对坐标解码回像素坐标、NMS 去重。

def decode_predictions(raw_output, conf_threshold=0.25, iou_threshold=0.4): # raw_output: (7, 7, 30),已经经过 sigmoid 处理过的概率部分 boxes = raw_output[..., :10].reshape(7, 7, 2, 5) class_probs = raw_output[..., 10:] class_probs = tf.sigmoid(class_probs) out_boxes = [] out_scores = [] out_classes = [] for i in range(7): for j in range(7): for b in range(2): x, y, w, h, conf = boxes[i, j, b] cls_prob = class_probs[i, j] * conf score = tf.reduce_max(cls_prob) cls_id = tf.argmax(cls_prob) if score < conf_threshold: continue # 解码回 448 尺寸下的绝对坐标 cx = (x + j) * 64 cy = (y + i) * 64 x1 = cx - w * 448 / 2 y1 = cy - h * 448 / 2 x2 = cx + w * 448 / 2 y2 = cy + h * 448 / 2 out_boxes.append([x1, y1, x2, y2]) out_scores.append(score.numpy()) out_classes.append(cls_id.numpy()) if len(out_boxes) == 0: return [], [], [] selected = tf.image.non_max_suppression( out_boxes, out_scores, max_output_size=100, iou_threshold=iou_threshold) return [out_boxes[i] for i in selected.numpy()], \ [out_scores[i] for i in selected.numpy()], \ [out_classes[i] for i in selected.numpy()]

NMS 的输入框要用像素坐标,不能用归一化或网格相对坐标。我见过有人把网格相对坐标直接丢进 NMS,导致不同网格的框因为坐标基准不同,IoU 算出来全是 0,去重完全失效,最后一张图上框叠成山。tf.image.non_max_suppression的max_output_size设成 100 足够覆盖单张图的目标数量,设太大也没有意义,只会拖慢后处理。

5.3 避坑清单:五条实际踩过的记录

现象一:训练 loss 一直不降,稳定在某个高值附近。原因是标签编码里中心点坐标没有做网格偏移,网络学习的 x_cell 和 y_cell 是绝对坐标,梯度永远对不齐。解决方法是回到encode_label,确认x_cell = cx * grid - grid_x这一步存在,并打印几个样本的 label 张量肉眼比对。

现象二:验证集 mAP 看着不错,但检测框整体偏移一个网格。原因是推理时解码坐标写成了cx = x + (j + 1)或cx = x * 64 + j * 64。x 已经是相对网格的比例,要先乘 64 得

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询