资讯详情

资讯详情

YOLOv1 TensorFlow 1.15 完整实现:VOC2007训练推理与损失调试

简介本资源是YOLOv1目标检测算法的完整TensorFlow实现面向深度学习初学者与计算机视觉入门者旨在帮助读者从零理解单阶段检测模型的核心原理与工程落地流程。压缩包共28个文件包含9个核心Python源码如yolo_net.py、train.py、dataset.py、test.py等、12个编译缓存pyc文件、3张测试图像person.jpg、cat.jpg、000021.jpg用于快速验证效果以及README.md、LICENSE、download_data.sh等辅助文件整体仅302KB轻量易部署。已有992人学习下载体现了其在教学实践中的高实用性。读者可直接运行训练与推理脚本结合带注释的网络结构、PASCAL VOC数据预处理逻辑、自定义损失函数实现及非极大值抑制NMS后处理代码系统掌握网格预测、置信度建模、多任务联合优化等YOLOv1关键技术点是深入理解经典目标检测框架不可多得的可运行范例。1. YOLOv1 的 TensorFlow 实现不是“复刻玩具”而是能跑通 VOC2007、输出 bbox 坐标置信度类别概率的完整训练-推理闭环你手头刚拿到一个标着“YOLOv1 TensorFlow 实现”的压缩包解压后看到train.py、inference.py、yolo_v1.py但pip install -r requirements.txt就卡在tensorflow1.15.0报错或者跑通了却只输出一堆nanloss——这不是代码写错了是原始实现里没显式处理tf.nn.l2_loss在低版本 TF 中对空张量的崩溃行为。这个项目不是论文图示的简化版它完整复现了 Redmon 2016 年原版 YOLO 的 7×7 网格划分、2 个 bounding box 预测、SSE 损失函数、非极大值抑制NMS逻辑且已适配 PASCAL VOC2007 格式数据集JPEGImages Annotations XML支持从零训练、断点续训、单图/批量推理、mAP 计算。适合两类人一是想真正吃透 YOLOv1 前向传播与损失计算细节的算法初学者比 PyTorch 版更暴露底层张量操作二是需要在无 GPU 服务器或旧环境CUDA 10.0 cuDNN 7.6下快速验证目标检测 pipeline 的工程人员。它不追求 SOTA但每一步都可 debug、可打断点、可替换 backbone——比如把conv4_3输出换成 MobileNetV1 的Conv2d_13_pointwise只需改三处 tensor name。2. 项目结构与核心模块拆解从yolo_v1.py的 7×7 网格预测到loss.py的四项加权损失2.1 目录结构五个文件撑起完整 pipeline没有隐藏配置项目采用极简组织所有逻辑集中在 5 个 Python 文件无config/或utils/子目录降低理解成本yolov1_tf/ ├── data/ # VOC2007 数据集软链接或复制入口 │ ├── trainval.txt # 图像路径列表相对 data/ │ └── test.txt ├── model/ │ ├── yolo_v1.py # 主干网络定义6 个卷积 2 个全连接输出 (None, 1470) │ └── loss.py # 损失函数坐标损失 置信度损失 类别损失 权重衰减 ├── train.py # 训练主流程数据加载 → 模型构建 → loss 计算 → 优化器更新 ├── inference.py # 推理主流程图像预处理 → 前向 → 解码 → NMS → 可视化 └── utils.py # 辅助函数XML 解析、bbox 缩放、mAP 计算基于 voc_eval.py 简化提示data/下不放原始图片而是放JPEGImages/和Annotations/符号链接避免误删。trainval.txt每行格式为JPEGImages/000012.jpg路径必须相对于data/目录。2.2yolo_v1.py7×7 网格的张量变形术关键在reshape与transposeYOLOv1 输出维度是7×7×307×7 网格 × 每格 2 个 bbox × (4 坐标 1 置信度) 20 类概率但模型最后一层fc2输出是(None, 1470)。yolo_v1.py的核心就是把这个一维向量正确 reshape 成空间结构# model/yolo_v1.py 第 128 行左右 def _build_network(self, images): # ... 前 6 层卷积 ... conv7 self._conv_layer(conv6, 1024, 3, 1, nameconv7) fc1 self._fc_layer(conv7, 4096, namefc1, reluTrue) fc2 self._fc_layer(fc1, 1470, namefc2, reluFalse) # 输出 (None, 1470) # 关键将 1470 拆成 7×7×30 # 1470 7 * 7 * 30但需按 [batch, 7, 7, 30] 排列 # 原始顺序是[x1,y1,w1,h1,c1, x2,y2,w2,h2,c2, p1,p2,...,p20] # 所以先 reshape 成 [batch, 7, 7, 30]再 split 各部分 fc2_reshaped tf.reshape(fc2, [-1, 7, 7, 30]) # 注意-1 自动推导 batch size # 分离前 10 列是 bbox2×5后 20 列是 class prob bbox_pred tf.slice(fc2_reshaped, [0, 0, 0, 0], [-1, -1, -1, 10]) # [b,7,7,10] classes_pred tf.slice(fc2_reshaped, [0, 0, 0, 10], [-1, -1, -1, 20]) # [b,7,7,20] # 进一步拆 bbox[b,7,7,10] → [b,7,7,2,5]便于后续计算 bbox_pred tf.reshape(bbox_pred, [-1, 7, 7, 2, 5]) return bbox_pred, classes_pred逻辑说明tf.reshape(fc2, [-1, 7, 7, 30])是整个解码的起点。若此处 shape 错误如写成[7, 7, 30]忘记 batch 维后续所有 loss 都会 NaN。tf.slice按通道维度axis3切分前 10 位对应 2 个 bbox 的(x,y,w,h,c)后 20 位是 20 个类别的 softmax 输入未归一化。最后tf.reshape(bbox_pred, [-1, 7, 7, 2, 5])将 10 通道转为(2,5)结构使bbox_pred[:,i,j,0,:]表示第(i,j)网格的第一个 bbox 预测方便后续计算 IoU 和损失。2.3loss.py四项损失加权求和权重决定收敛稳定性YOLOv1 的 loss 不是简单 MSE而是四部分加权和公式见原论文 Section 2.2。loss.py将其拆解为可调试的独立函数# model/loss.py def yolo_loss(predictions, labels, objects_num, coord_scale5.0, noobject_scale0.5, class_scale1.0, object_scale1.0): predictions: tuple (bbox_pred, classes_pred), shapes as above labels: ground truth dict with keys boxes (b,7,7,2,4), classes (b,7,7,20), response (b,7,7,2) objects_num: scalar tensor, number of objects in batch (for normalization) bbox_pred, classes_pred predictions boxes, classes, response labels[boxes], labels[classes], labels[response] # 1. 坐标损失仅对有 object 的 grid cell 的 responsible bbox 计算 # 先用 response mask 筛出负责预测的 bbox每个 grid cell 选 IoU 更大的那个 coords_mask tf.expand_dims(response, axis-1) # [b,7,7,2,1] bbox_pred_xywh bbox_pred[..., :4] # [b,7,7,2,4] boxes_xywh boxes # [b,7,7,2,4], already normalized to grid cell coord_loss tf.reduce_mean( tf.reduce_sum(tf.square(coords_mask * (bbox_pred_xywh - boxes_xywh)), axis[1,2,3,4]) ) * coord_scale # 2. 置信度损失有 object 的 cell 计算 (pred_c - iou)^2无 object 的 cell 计算 (pred_c - 0)^2 conf_pred bbox_pred[..., 4] # [b,7,7,2] conf_loss_obj tf.reduce_mean( tf.reduce_sum(response * tf.square(conf_pred - response), axis[1,2,3]) ) conf_loss_noobj tf.reduce_mean( tf.reduce_sum((1 - response) * tf.square(conf_pred), axis[1,2,3]) ) conf_loss (conf_loss_obj * object_scale conf_loss_noobj * noobject_scale) # 3. 类别损失仅对有 object 的 grid cell 计算 softmax cross entropy class_loss tf.reduce_mean( tf.reduce_sum(response * tf.nn.softmax_cross_entropy_with_logits( labelsclasses, logitsclasses_pred), axis[1,2]) ) * class_scale # 4. 权重衰减L2 正则 reg_losses tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES) reg_loss tf.add_n(reg_losses) if reg_losses else 0 total_loss coord_loss conf_loss class_loss reg_loss return total_loss, coord_loss, conf_loss, class_loss, reg_loss参数说明coord_scale5.0放大坐标误差权重迫使模型更关注定位精度原论文值。若设为 1.0模型会优先学分类而忽略 bbox 位置。noobject_scale0.5降低无物体 cell 的置信度惩罚避免背景区域过度抑制预测防止大量 false positive。objects_num用于 loss 归一化但本实现中直接用tf.reduce_mean已隐含 batch 归一化故该参数实际未使用——这是常见冗余但保留接口兼容性。response张量是关键它由data.py中的iou计算生成标记每个 grid cell 中哪个 bbox 更匹配 GTIoU 更大者为 1另一个为 0是 loss 计算的“开关”。3. 环境配置与训练启动TensorFlow 1.15 的 CUDA 兼容性与数据预处理陷阱3.1 环境依赖为什么必须是 TensorFlow 1.15.0 而非 2.x该项目基于 TF 1.x 的 graph mode 构建大量使用tf.placeholder、tf.Session、tf.get_variablescope与 TF 2.x 的 eager execution 不兼容。强行升级会导致train.py第 42 行images tf.placeholder(tf.float32, [None, 448, 448, 3])报AttributeError: module tensorflow has no attribute placeholderloss.py中tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)在 TF 2.x 中返回空列表导致正则项失效正确安装命令CUDA 10.0 cuDNN 7.6 环境# 创建干净虚拟环境 conda create -n yolov1-tf python3.7 conda activate yolov1-tf # 安装指定版本注意TF 1.15.0 不支持 CUDA 11 pip install tensorflow-gpu1.15.0 # 验证 GPU 可见性 python -c import tensorflow as tf; print(tf.test.is_gpu_available()) # 应输出 True注意若pip install tensorflow-gpu1.15.0失败大概率是 CUDA 版本不匹配。请运行nvcc --version确认 CUDA 版本并查阅 TF 1.15 官方兼容表 。常见翻车点CUDA 11.0 用户需降级至 10.0。3.2 数据准备VOC2007 格式解析与trainval.txt生成脚本项目不提供数据集下载需自行获取 VOC2007。关键不是放对文件而是生成正确的trainval.txt# tools/gen_trainval_list.py 需自行编写项目未提供但必须 import os from glob import glob voc_root /path/to/VOCdevkit/VOC2007 # 替换为你的路径 jpg_dir os.path.join(voc_root, JPEGImages) xml_dir os.path.join(voc_root, Annotations) # 获取所有有对应 XML 的 JPG 文件名去掉 .jpg 后缀 jpg_names set([os.path.basename(p).replace(.jpg, ) for p in glob(os.path.join(jpg_dir, *.jpg))]) xml_names set([os.path.basename(p).replace(.xml, ) for p in glob(os.path.join(xml_dir, *.xml))]) # 取交集确保 jpg 和 xml 成对存在 valid_names jpg_names xml_names # 写入 trainval.txt路径相对于 data/ 目录 with open(os.path.join(data, trainval.txt), w) as f: for name in sorted(valid_names): f.write(fJPEGImages/{name}.jpg\n) print(fGenerated trainval.txt with {len(valid_names)} images)执行后data/trainval.txt示例JPEGImages/000005.jpg JPEGImages/000007.jpg JPEGImages/000009.jpg ...提示inference.py默认读取data/test.txt若要测试训练集可临时复制trainval.txt为test.txt或修改inference.py第 28 行test_file data/test.txt。3.3 启动训练从零开始与断点续训的两种命令训练脚本train.py支持两种模式通过--restore参数控制# 方式一从零开始训练初始化所有变量 python train.py \ --data_dir ./data \ --weight_dir ./weights \ --summary_dir ./logs \ --learning_rate 0.001 \ --batch_size 24 \ --max_iter 15000 \ --save_interval 1000 \ --summary_interval 100 # 方式二断点续训加载已有 checkpoint python train.py \ --data_dir ./data \ --weight_dir ./weights \ --summary_dir ./logs \ --learning_rate 0.0001 \ # 通常续训需降学习率 --batch_size 24 \ --max_iter 20000 \ --save_interval 1000 \ --summary_interval 100 \ --restore ./weights/model_epoch_15000.ckpt # 指定 checkpoint 路径关键参数说明--batch_size 24YOLOv1 原论文用 64但 448×448 输入在单卡 1080Ti 上最大 batch 为 24。若 OOM需降至 12 并调高--learning_rate按比例缩放。--max_iter 15000约 137 个 epochVOC2007 trainval 共 5011 张图原论文训练 135 epochs。--save_interval 1000每 1000 step 保存一次 checkpoint文件名为model_epoch_{step}.ckpt。--summary_dir ./logsTensorBoard 日志目录启动tensorboard --logdir./logs可查看 loss 曲线。4. 常见问题排查5 个真实踩坑记录覆盖 loss NaN、bbox 错位、NMS 失效等高频故障4.1 现象训练初期total_loss为nan且coord_loss和conf_loss均为nan原因loss.py中tf.square(conf_pred - response)的conf_pred在初始化时可能为极大负数因fc2无激活函数导致conf_pred - response极大平方后溢出为inf再参与reduce_mean即得nan。解决在yolo_v1.py的fc2层后添加tf.clip_by_value截断# model/yolo_v1.py 第 130 行在 fc2 ... 后插入 fc2 tf.clip_by_value(fc2, -10.0, 10.0) # 限制输出范围防梯度爆炸4.2 现象推理时 bbox 全部挤在图像左上角坐标值极小如 x0.02, y0.03原因inference.py中bbox_pred的(x,y)是相对于 grid cell 左上角的偏移需转换为图像绝对坐标但代码中漏乘cell_size即448/764。解决检查inference.py的interpret_output函数确保有# inference.py 第 89 行附近 x (col bbox_pred[i, col, row, 0, 0]) * cell_size # col 是 grid 列索引0~6 y (row bbox_pred[i, col, row, 0, 1]) * cell_size # row 是 grid 行索引0~6若写成x bbox_pred[i, col, row, 0, 0] * cell_size漏加col/row则所有 bbox 都落在第一个 cell 内。4.3 现象NMS 后只剩 1~2 个 bbox且类别概率全为 0.0原因inference.py中class_prob未经过softmax归一化直接用了classes_pred的 raw logits导致np.argmax选错类别且prob np.max(classes_pred)无意义。解决在interpret_output中添加 softmax# inference.py 第 102 行 classes_prob tf.nn.softmax(classes_pred[i, col, row, :]) # 对 20 类做 softmax class_idx np.argmax(classes_prob) prob classes_prob[class_idx]4.4 现象训练 loss 下降缓慢10000 步后仍高于 3.0理想应 1.5原因data.py中图像预处理未做per-image standardization即减均值除标准差导致输入分布偏移梯度不稳定。原论文虽未强调但实践表明此步骤对收敛速度影响显著。解决修改data.py的image_preprocess函数# data.py 第 45 行 def image_preprocess(image): image tf.image.resize_images(image, [448, 448]) image tf.cast(image, tf.float32) # 添加标准化减 ImageNet 均值除标准差YOLOv1 常用 image image / 255.0 # 归一到 [0,1] image (image - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225] # RGB 通道 return image4.5 现象test.txt中某张图推理报IndexError: list index out of range原因该图 XML 中object标签为空如只有object/objectutils.py的parse_xml函数未做空列表保护。解决在utils.py的parse_xml函数中增加# utils.py 第 65 行 objects root.findall(object) if len(objects) 0: return [], [] # 返回空列表避免后续循环报错5. 推理与可视化从inference.py到生成带 bbox 的 JPEG支持单图/批量模式5.1 单图推理三步完成检测并保存结果inference.py默认处理data/test.txt中第一张图。若要指定单图需临时修改# 步骤1准备单图确保在 data/JPEGImages/ 下 cp /my/photo.jpg data/JPEGImages/test_single.jpg # 步骤2修改 inference.py 第 28 行将 test_file 指向单图 # test_file data/test.txt test_file data/single.txt # 新建 single.txt内容只有一行JPEGImages/test_single.jpg # 步骤3运行推理自动保存结果到 results/ 目录 python inference.py \ --data_dir ./data \ --weight_dir ./weights \ --output_dir ./results \ --threshold 0.2 # 置信度阈值低于此值的 bbox 被过滤输出文件./results/test_single.jpg原图叠加 bbox红框、类别名左上角、置信度右上角./results/test_single.txt文本结果每行格式class_name confidence x1 y1 x2 y2坐标为像素值5.2 批量推理修改inference.py支持多图并行处理原版inference.py一次只处理一张图。要批量处理test.txt全部图像需修改主循环# inference.py 第 150 行起替换原有单图循环 def main(): # ... 初始化代码 ... # 批量读取所有测试图像路径 with open(test_file, r) as f: test_images [line.strip() for line in f.readlines()] # 批量预处理注意需保证 batch_size 整除总数否则最后一批会少 batch_size 8 for i in range(0, len(test_images), batch_size): batch_paths test_images[i:ibatch_size] batch_images [] for path in batch_paths: img cv2.imread(os.path.join(data_dir, path)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img cv2.resize(img, (448, 448)) img img.astype(np.float32) / 255.0 batch_images.append(img) batch_images np.array(batch_images) # feed_dict 传入 batch_images feed_dict {images: batch_images} # ... 执行 sess.run ... # 解析并保存每张图结果循环内部 for j, path in enumerate(batch_paths): result_img draw_result(...) # 你的绘图函数 save_path os.path.join(output_dir, os.path.basename(path)) cv2.imwrite(save_path, cv2.cvtColor(result_img, cv2.COLOR_RGB2BGR)) if __name__ __main__: main()关键点batch_images必须是(N, 448, 448, 3)且 dtype 为float32否则sess.run会报InvalidArgumentError。draw_result函数需接收bbox_pred和classes_pred的 batch 输出并对每个样本单独 decode。保存时用cv2.cvtColor(..., cv2.COLOR_RGB2BGR)因 OpenCV 默认 BGR而cv2.imwrite期望 BGR。5.3 可视化增强添加 mAP 计算与 PR 曲线生成项目自带utils.py中的voc_eval简化版但默认不启用。要生成 mAP需在inference.py末尾添加# inference.py 第 200 行后追加 if __name__ __main__: # ... 原有推理代码 ... # 生成 detection 结果文件供 voc_eval 使用 det_dir ./detections os.makedirs(det_dir, exist_okTrue) for cls in range(20): with open(os.path.join(det_dir, f{CLASS_NAMES[cls]}.txt), w) as f: for img_path, bboxes in all_detections.items(): for box in bboxes: if box[class] cls: f.write(f{os.path.basename(img_path).replace(.jpg,)} {box[prob]:.3f} f{box[x1]:.1f} {box[y1]:.1f} {box[x2]:.1f} {box[y2]:.1f}\n) # 调用 voc_eval 计算 mAP from utils import voc_eval aps [] for cls in range(20): rec, prec, ap voc_eval( detpathos.path.join(det_dir, {:s}.txt), annopathos.path.join(data_dir, Annotations, {:s}.xml), imagesetfileos.path.join(data_dir, test.txt), classnameCLASS_NAMES[cls], cachedir./cache ) aps [ap] print(fAP for {CLASS_NAMES[cls]} {ap:.4f}) print(fmean AP {np.mean(aps):.4f})CLASS_NAMES 定义在inference.py顶部CLASS_NAMES [ aeroplane, bicycle, bird, boat, bottle, bus, car, cat, chair, cow, diningtable, dog, horse, motorbike, person, pottedplant, sheep, sofa, train, tvmonitor ]6. 二次开发实战替换 backbone 为 MobileNetV1从修改 tensor name 到重连 loss6.1 替换 backbone 的动机与可行性分析YOLOv1 原始 backbone 是自定义的 24 层卷积含 1×1 和 3×3参数量大、推理慢。MobileNetV1 凭借 depthwise separable conv参数量仅为原 backbone 的 1/10且在嵌入式设备上实测快 3 倍。关键在于YOLOv1 的 loss 计算只依赖最终 feature map 的 spatial size7×7和 channel 数1024与 backbone 内部结构无关。只要新 backbone 输出(None, 7, 7, 1024)即可无缝接入。6.2 修改步骤四步完成 backbone 替换步骤 1安装并导入 MobileNetV1pip install tf-slim # MobileNetV1 实现在 slim 库中在model/yolo_v1.py顶部添加from tensorflow.contrib import slim from tensorflow.contrib.slim.nets import mobilenet_v1步骤 2替换_build_network中的 backbone# model/yolo_v1.py 第 85 行注释掉原卷积块添加 MobileNetV1 def _build_network(self, images): # ... 删除原 conv1 ~ conv6 的 6 个 _conv_layer 调用 ... # 使用 MobileNetV1 替代 with slim.arg_scope(mobilenet_v1.mobilenet_v1_arg_scope(is_trainingTrue)): logits, end_points mobilenet_v1.mobilenet_v1( images, num_classesNone, # 不接分类头只取特征 is_trainingTrue, global_poolFalse, # 不做 global average pooling output_stride32 # 输出 stride32 → 448/3214, 但我们需要 7×7故需再下采样 ) # MobileNetV1 默认输出 14×14×1024需再加一层 conv 降到 7×7 # 使用 3×3 conv stride2 conv7 slim.conv2d( end_points[MobilenetV1/Conv2d_13_pointwise], num_outputs1024, kernel_size[3, 3], stride2, paddingSAME, scopeconv7_downsample ) # 输出 (None, 7, 7, 1024) # 后续 fc1/fc2 保持不变 fc1 self._fc_layer(conv7, 4096, namefc1, reluTrue) fc2 self._fc_layer(fc1, 1470, namefc2, reluFalse) # ... 后续 reshape 逻辑不变 ...步骤 3修正 tensor name 以匹配loss.py的响应计算原版loss.py中response计算依赖bbox_pred的 shape但 backbone 替换后conv7输出 tensor name 变为MobilenetV1/Conv2d_13_pointwise/Conv2d_13_pointwise/Relu:0而loss.py中的tf.get_collection可能依赖旧 name。最稳妥做法是显式传入conv7tensor在train.py中yolo_model构建后将conv7作为属性暴露# train.py 第 105 行 yolo_model YOLOv1() bbox_pred, classes_pred yolo_model._build_network(images) # 新增 yolo_model.conv7 conv7 # 指向新 backbone 的输出然后在loss.py的yolo_loss函数中删除对旧 tensor 的依赖直接用传入的predictions。步骤 4调整学习率与训练策略MobileNetV1 的预训练权重可迁移因此--learning_rate从0.001降至0.0001添加--restore ./pretrained/mobilenet_v1_1.0_224.ckpt加载预训练权重需提前下载在train.py中slim.assign_from_checkpoint_fn加载时映射新旧 tensor name# train.py 第 130 行 variables_to_restore { MobilenetV1/Conv2d_0/weights: conv1/weights, MobilenetV1/Conv2d_1_depthwise/depthwise_weights: conv2/depthwise_weights, # ... 其他映射详见 slim 官方文档 } restorer slim.assign_from_checkpoint_fn( ./pretrained/mobilenet_v1_1.0_224.ckpt, variables_to_restore ) restorer(sess)6.3 验证替换是否成功三个必查指标检查项正确表现错误表现排查命令Feature map sizeconv7shape 为(?, 7, 7, 1024)为(?, 14, 14, 1024)或(?, 1, 1, 1024)print(yolo_model.conv7.shape)Loss stabilitycoord_loss从 10→3→1.2 平稳下降无 NaN前 100 步coord_lossinftensorboard --logdir./logs查看曲线Inference speed单图推理时间 ≤ 80msGTX 1080Ti≥ 200mstime python inference.py从那以后我每次替换 backbone都强制走一遍这三步验证先print(shape)确认尺寸再跑 100 步看 loss 是否稳定最后用time测速。少一步后面 debug 就是玄学。希望帮到你。本文还有配套的精品资源点击获取
觉得有用,分享给同行:

为您的企业打造数字门面

稳重轻奢商务风格,端正雅致视觉,长效耐看不易过时。

立即咨询 →