用TensorFlow手写YOLOv1目标检测:从标签编码到损失函数与推理避坑

发布时间:2026/10/12 21:42:13
用TensorFlow手写YOLOv1目标检测:从标签编码到损失函数与推理避坑 简介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×520每个框有中心坐标 x、y宽高 w、h和一个置信度然后整格共享 20 个类别的条件概率。先记住这个编码后面所有代码都在围绕它转。标签编码的第一步是归一化。所有坐标都以整张图宽高为基准归一化到 0~1中心点还要进一步转成「相对于所在网格」的坐标。宽高保持相对整图的归一化值这样在 0~1 之间更重要的是训练时避免大图小图比例不一致导致回归不稳定。这一步一定要归一化不然损失函数里的坐标项会被像素量大数值带偏整个训练直接失控。第二个关键点是「谁负责这个目标」。每个物体只会落到一个网格里由中心点所在网格负责该网格的两个预测框里训练时会计算两个框和真实框的 IoUIoU 更大的那个被标记为 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), dtypenp.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_listtf.train.BytesList(value[img_bytes])), label: tf.train.Feature(bytes_listtf.train.BytesList(value[label_bytes])), } example tf.train.Example(featurestf.train.Features(featurefeature)) 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], channels3) 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_delta0.05) image tf.image.random_saturation(image, lower0.8, upper1.2) return image, labelflip_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, stride1): x layers.Conv2D(filters, kernel, stridesstride, paddingsame)(x) x layers.BatchNormalization()(x) return layers.LeakyReLU(alpha0.1)(x) def build_yolov1(num_grid7, num_boxes2, num_classes20): inputs tf.keras.Input(shape(448, 448, 3)) x conv_block(inputs, 64, 7, stride2) 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, stride2) x conv_block(x, 1024, 3) x layers.Flatten()(x) x layers.Dense(4096)(x) x layers.LeakyReLU(alpha0.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 之后的值。你在实现时如果手滑给最后一层加了activationsigmoid训练初期看起来 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 是这个映射后的权重 dictkey 为层名 for layer in model.layers: if layer.name in pretrained: layer.set_weights(pretrained[layer.name]) print(floaded {layer.name})如果你只想快速跑通完整流程不追求和原版完全一致的 mAP可以直接用上面的结构从头训练把学习率调小到 1e-4batch size 用 16 或 32在小数据集上跑几十轮也能看到框逐渐变准。但要做正经对比实验、要复现论文指标预训练权重几乎是必须的。4. 损失函数YOLOv1 最核心的五项设计4.1 损失拆解坐标、宽高、置信度、类别权重不一样YOLOv1 的损失函数是五项之和每一项都带着自己的权重系数。坐标误差和宽高误差权重是 5因为坐标直接决定框的位置必须优先学准不含目标的置信度误差权重是 0.5因为绝大多数网格里根本没有目标如果不压低权重网络会学成「永远预测没有目标」的偷懒解。宽高误差用的是平方根差。这个设计的直觉是同样差 10 个像素对小框影响远比大框严重开平方后让小框的误差在数值上被放大梯度也会更大网络就会更努力去修正小框。这在今天看有点粗糙但在当时是对「大小框公平性」的一次有效探索。置信度部分要注意 y_pred 的置信度是线性输出值y_true 的置信度则取「预测框和真实框的 IoU」。就是说训练目标不是简单设为 1而是这个框自己和目标的实际重合程度。这样网络学出来的置信度天然带有定位质量信息推理时可以直接用来做 NMS 排序。类别损失只在一个网格内有目标时计算而且类别概率是网格共享的和两个框无关。这就是 YOLOv1 对密集目标无能为力的结构性原因损失函数层面没有任何一个位置能表达「同一个格子有两个不同类」。4.2 TensorFlow 实现mask 是怎么生成的损失实现的核心是生成三个 maskresponsible mask、noobj mask、object mask。responsible mask 决定哪个预测框参与坐标和置信度回归noobj mask 决定哪些框要被压到 0 置信度。def yolo_loss(y_true, y_pred, grid7, boxes2, classes20, lambda_coord5.0, lambda_noobj0.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, keepdimsTrue) 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_rate1e-5end_learning_rate1e-4再配合decay_steps覆盖前 10 轮。直接手写回调控制optimizer.lr也可以但容易在 TF2 的图执行模式下失效。5.2 推理阶段解码输出再 NMS别直接用网络输出画框训练结束后拿到的模型输出还是 7×7×30 的原始张量直接画框会有大量重叠框。因为每个目标附近会有多个网格同时预测到它而且同一网格的两个框可能都在往目标上靠。推理时要先做三件事过滤低置信度、把框坐标从网格相对坐标解码回像素坐标、NMS 去重。def decode_predictions(raw_output, conf_threshold0.25, iou_threshold0.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_size100, iou_thresholdiou_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 得本文还有配套的精品资源点击获取

关于本文作者

来自尧图内容编辑团队

尧图内容编辑团队 内容团队

尧图内容编辑团队

本文由尧图网络内容编辑团队执笔。团队由资深项目经理、前端工程师与设计师组成,所有内容均来自亲手交付的真实项目,先讲清问题、再给出可落地的解法。尧图深耕北京网站建设十年,服务过京华建材集团、智造科技等各行业客户,把一线经验沉淀为可复用的行业观察。

  • 十年建站经验,覆盖建材、制造、服务、文创等
  • 项目经理把关选题与事实准确性
  • 工程师与设计师联合撰写专业细节
  • 统一编辑规范,保证文风与排版一致
  • 每月复盘转化数据,迭代选题方向

延伸阅读

相关资讯与近期热门内容

深度阅读推荐

建站决策前值得细读的三篇

网站改版的5个关键决策
2024-08-12

网站改版的5个关键决策

什么时候该改版、改到什么程度、如何避免流量掉光,京华建材集团改版复盘给出答案。

获取专属建站方案

看完文章,把您的行业与预算告诉我们,免费获取一份量身定制的官网建设方案与报价。

立即免费咨询