ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

PyTorch实现YOLOv3-tiny:从Darknet权重加载到边缘部署

PyTorch实现YOLOv3-tiny:从Darknet权重加载到边缘部署 简介一份面向边缘设备与实时检测场景的轻量级目标检测实现基于PyTorch构建完整的YOLOv3-tiny开发流程。作为简化版YOLOv3模型通过减少网络层数在精度与速度之间取得平衡特别适合硬件资源受限的嵌入式或移动端应用。整个包体共22个文件以14个Python脚本为核心覆盖模型定义、锚点计算、LMDB数据构建、预训练微调与推理预测等环节另含类别标签、网络结构示意图和说明文档压缩包约1.17MB便于快速下载与部署。资源面向需要入门目标检测或优化边缘端推理性能的开发者配有可运行的测试脚本与示例图片从卷积层设计到损失函数从锚点机制到数据加载能够帮助深入理解YOLOv3-tiny的完整实现细节。在此基础上可针对自定义数据集调整参数训练适配自身需求的实时检测器。当前已有49人学习下载适合结合配套说明逐步对照实践。1. 为什么还要自己写一个 YOLOv3-tiny边缘设备上的实时检测选型我在某智能硬件项目上被帧率卡了三天最后又把 YOLOv3-tiny 捡了回来。当时团队主推的方案在开发机上跑得飞快上到目标设备后直接掉到个位数帧换成 tiny 才把实时检测的底线保住。这个 zip 标题背后的东西就是一套能看懂、能改动、能自己训练的 YOLOv3-tiny 的 PyTorch 实现。对从业者来说它解决三件事官方权重文件里装的到底是什么、模型怎么在自己数据上微调、以及部署时哪些环节最容易翻车。适合手里有边缘设备、正在做实时目标检测选型的人也适合刚入门想亲手拆一个检测模型的开发者。2. 搭网络PyTorch 从零实现 YOLOv3-tiny 的 backbone 与双尺度检测头2.1 先看结构表tiny 为什么只有七层卷积还能检测YOLOv3-tiny 的核心卖点是轻。它把 Darknet-53 换成了一套只有七层卷积的 backbone再靠两个检测头撑起多尺度检测。做边缘部署的人选它不是因为精度高而是因为计算量小、显存占用低、推理延迟可控在树莓派和嵌入式设备这类算力有限的平台上这是比 YOLOv5s 更稳的底线方案。网络结构可以用一张表说清楚输入是 416x416x3经过下采样后分出两条支路模块主要算子输出尺寸用途backbone5 次卷积 4 次 maxpool26x26x256提取浅层特征也是第二检测头的输入det_13maxpool 保持尺寸 5 组卷积13x13x255负责大目标检测route 支路上采样 concat 2 组卷积26x26x255负责小目标检测注意最后的 255 是怎么来的3 个 anchor 乘以 8585 等于 xywh 四项加一个物体置信度加 80 类置信度。如果你换一个类别的数据集最后这个 255 要改成 3 乘5 加类别数。很多人只改前面 backbone 不动尾部训练时 loss 直接 NaN多半就是这里出了问题。2.2 用 nn.Module 把网络写出来下面这段代码是我平时实现 YOLOv3-tiny 的写法把 backbone、13x13 检测头、26x26 检测头拆成三个子模块forward 里只保留一个 route 节点逻辑比用单一顺序列表更清楚import torch import torch.nn as nn import torch.nn.functional as F def conv_bn_lrelu(in_c, out_c, k, s1, p1, bnTrue): layers [nn.Conv2d(in_c, out_c, k, strides, paddingp, biasnot bn)] if bn: layers.append(nn.BatchNorm2d(out_c)) layers.append(nn.LeakyReLU(0.1, inplaceTrue)) return layers class YOLOv3Tiny(nn.Module): def __init__(self, num_classes80): super().__init__() self.num_classes num_classes # backbone输出 26x26x256 self.backbone nn.Sequential( *conv_bn_lrelu(3, 16, 3), nn.MaxPool2d(2, 2), *conv_bn_lrelu(16, 32, 3), nn.MaxPool2d(2, 2), *conv_bn_lrelu(32, 64, 3), nn.MaxPool2d(2, 2), *conv_bn_lrelu(64, 128, 3), nn.MaxPool2d(2, 2), *conv_bn_lrelu(128, 256, 3), ) # 13x13 检测头 self.maxpool_keep nn.MaxPool2d(2, 1, padding1) self.conv512 nn.Sequential(*conv_bn_lrelu(256, 512, 3)) self.conv1024 nn.Sequential(*conv_bn_lrelu(512, 1024, 3)) self.conv256_1 nn.Sequential(*conv_bn_lrelu(1024, 256, 1, p0)) self.conv512_2 nn.Sequential(*conv_bn_lrelu(256, 512, 3)) self.det13_conv nn.Conv2d(512, 3 * (num_classes 5), 1) # 26x26 检测头 self.upsample nn.Upsample(scale_factor2, modenearest) self.conv256_2 nn.Sequential(*conv_bn_lrelu(1024 256, 256, 3)) self.det26_conv nn.Conv2d(256, 3 * (num_classes 5), 1) def forward(self, x): feat_26 self.backbone(x) # 26x26x256 x13 self.maxpool_keep(feat_26) x13 self.conv512(x13) x13 self.conv1024(x13) route_1024 x13 # 给第二检测头做 route x13 self.conv256_1(x13) x13 self.conv512_2(x13) det_13 self.det13_conv(x13) # 13x13x255 x26 self.upsample(route_1024) # 26x26x1024 x26 torch.cat((x26, feat_26), dim1) # concat x26 self.conv256_2(x26) det_26 self.det26_conv(x26) # 26x26x255 return det_13, det_26 if __name__ __main__: model YOLOv3Tiny(num_classes80) x torch.randn(1, 3, 416, 416) out13, out26 model(x) print(out13.shape, out26.shape)这段代码的 forward 顺序和官方 cfg 是严格对应的。关键点有三个第一个是maxpool_keep用的是 stride1、padding1它不会改变 13x13 的尺寸很多复现版本在这里直接写MaxPool2d(2, 2)会导致后面的卷积全部在 12x12 上跑最终输出 shape 对不上权重文件第二个是 route 节点取的是 conv1024 的输出不是 conv256_1 之后的特征方向反了会丢掉语义信息小目标检测明显变差第三个是 26x26 那个检测头的输入是 1024 加 256 做 concat1024 来自上采样256 来自 backbone 浅层特征这个拼接顺序在加载权重、写导出脚本时都要保持一致。2.3 该用哪套初始化随机初始化、官方权重还是迁移学习搭完网络之后先别急着训练初始化方式决定了你是从零开始学还是站在前人肩膀上微调。如果你手头有已经标注好的私有数据集常见做法是加载 Darknet 官方 COCO 权重把最后的分类卷积层换成你自己的类别数再做迁移学习如果你纯粹是想把 YOLOv3-tiny 跑通、验证前向和 loss直接用 PyTorch 默认的 kaiming 初始化就够了。有个容易被忽略的点加载官方权重时最后一个检测头的输出通道数必须改。COCO 是 80 类255 通道如果你只有 2 类尾部要改成 3 乘 6 等于 18此时不能让加载器硬灌 255 通道的权重否则 shape 不匹配会直接报错。我一般会先把模型搭好打印每个张量的 shape再决定改哪些层避免在白板上推半天结构一跑就崩。3. 加载 Darknet 官方权重PyTorch 转换脚本与参数对位3.1 Darknet 权重文件里到底装了什么YOLOv3-tiny 的官方权重是一个二进制文件不是常见的 PyTorch 的 .pth 或 .ckpt。它从头到尾是一串 float32 的小端字节流最前面一般是几个 int32 的头部字段跨版本会有差异常见是 3 个或 4 个后面的数据严格按 cfg 里卷积层的顺序排列。每遇到一个带 batch_normalize 的卷积块先写 BN 的 bias、weight、running_mean、running_var最后写卷积核权重没有 BN 的层先写卷积 bias再写卷积核权重。为什么转换容易出错因为 PyTorch 的模型参数是分散挂在各个 module 上的而 Darknet 权重是一个扁平的数组。要么你手动按 cfg 顺序把参数一个个切出来要么写一个通用加载器遍历模型 module。手工硬编码层数对不了几个版本就会出问题我建议第一次就写通用加载器后面换数据集、改通道数都能复用。3.2 用 Python 写一个通用加载器下面这段加载器在我的几个模拟项目里反复用过核心思路是先解析 cfg 里每个卷积块有没有 BN再按顺序把二进制数组切块填进模型参数import numpy as np import torch import torch.nn as nn def parse_cfg_bn(cfg_path): bn_flags {} conv_idx 0 with open(cfg_path, r) as f: for line in f: line line.strip() if line.startswith([convolutional]): conv_idx 1 bn_flags[conv_idx] False elif line.startswith(batch_normalize): bn_flags[conv_idx] line.split()[1].strip() 1 return bn_flags def load_darknet_weights(model, cfg_path, weights_path): bn_flags parse_cfg_bn(cfg_path) with open(weights_path, rb) as fp: header np.fromfile(fp, dtypenp.int32, count3) w np.fromfile(fp, dtypenp.float32) idx 0 conv_idx 0 for m in model.modules(): if isinstance(m, nn.Conv2d): conv_idx 1 if not bn_flags.get(conv_idx, False): n m.bias.numel() m.bias.data.copy_(torch.from_numpy(w[idx: idx n])) idx n n m.weight.numel() m.weight.data.copy_( torch.from_numpy(w[idx: idx n]).view(m.weight.shape) ) idx n elif isinstance(m, nn.BatchNorm2d): for p in (m.bias, m.weight, m.running_mean, m.running_var): n p.numel() p.data.copy_(torch.from_numpy(w[idx: idx n])) idx n return idx, len(w)使用方式很简单构造好 YOLOv3Tiny 之后把 cfg 路径和 weights 路径传进去返回的idx和len(w)如果相等说明所有字节都被消费掉了。常见的错误是idx比len(w)小一截或者反过来读越界这通常是你模型里的卷积层顺序和 cfg 对不上也可能是模型尾部通道数改过之后参数总量变了。参数说明count3是读取头部三个 int32大部分官方权重适用如果你遇到加载后输出全 0先试count4或count5这是跨版本最常见的差异点。另外m.weight.data.copy_之前一定要做.view(m.weight.shape)因为 Darknet 卷积核存储顺序是[out_ch, in_ch, kh, kw]和 PyTorch 的[out_ch, in_ch, kh, kw]一致但直接 reshape 能防止目标数组长度对不上时出现静默错位。3.3 权重加载后的自检跑一张图对比检测结果加载完先别急着训练做一次快速自检。我用一张同时包含大目标和小目标的图片比如人物居中的街景图把模型切到 eval 模式跑一次前向打印两个检测头的输出张量均值、标准差和最大绝对值。如果输出张量出现大量 NaN 或全部是零不要怀疑模型写错先怀疑权重读错了。常见原因是头部字节数不对导致后续所有参数错位其次是模型里多了一个 PyTorch 自动初始化的卷积层而 cfg 里没有对应层导致 idx 越界。把加载器返回的idx打出来和len(w)对比能立刻定位问题。另一个更隐蔽的问题是 BN 参数顺序Darknet 是先写 bias、再写 weight、再写 mean、var我的加载器里for p in的顺序刚好对应这个顺序如果你自己写千万别写成 weight、bias、mean、var。4. 把张量变成检测框并跑通训练闭环4.1 batch 版坐标解码从 255 通道特征图还原出预测框模型输出的两个张量 shape 是[B, 255, 13, 13]和[B, 255, 26, 26]它们是未经 sigmoid 的 logits不能直接当框用。解码要做的就是把每个 grid cell 里的 3 个 anchor 预测拆成中心点、宽高、置信度和类别分数。这里有个细节Darknet 的原始输出在训练时会对置信度和类别做 sigmoid而 PyTorch 复现版一般让模型直接输出 logits到后处理再做避免重复计算。def decode_boxes(pred, anchors, stride, num_classes80): batch, ch, h, w pred.shape nA len(anchors) pred pred.view(batch, nA, num_classes 5, h, w) pred pred.permute(0, 1, 3, 4, 2).contiguous() grid_x, grid_y torch.meshgrid( torch.arange(w), torch.arange(h), indexingxy ) grid torch.stack((grid_x, grid_y), dim-1).float().to(pred.device) xy torch.sigmoid(pred[..., :2]) grid wh torch.exp(pred[..., 2:4]) * torch.tensor(anchors).float().to(pred.device) conf torch.sigmoid(pred[..., 4:5]) cls torch.sigmoid(pred[..., 5:]) boxes torch.cat( [(xy * stride), (wh * stride), conf, cls], dim-1 ) return boxes逻辑说明pred[..., :2]加上 grid 坐标后还要乘 stride是因为特征图每个 cell 对应原图上一个 stride 大小的区域13x13 特征图 stride 是 3226x26 是 16。wh用 exp 还原是因为网络学的是 log-space 的宽高偏移如果数据里目标尺度跨越特别大比如同时有几十像素的小目标和高清大图目标exp在训练初期容易产生极大值后续 loss 容易不稳定。参数说明anchors要按检测头分开传。13x13 用 mask 3、4、5对应(81,82)、(135,169)、(344,319)26x26 用 mask 0、1、2对应(10,14)、(23,27)、(37,58)。很多翻车现场就是把这两组 anchor 调换了表现是小目标一个都检不到大目标框又偏又大。4.2 接上 YOLOv3-tiny 的 Loss三类损失与正样本筛选训练用的 loss 主要由三部分构成坐标损失、置信度损失、类别损失。坐标损失可以简单用 MSE也可以换 CIoU但 YOLOv3 系列原版风格是分开算xy 用带 logits 的 MSE 或 BCEwh 用 MSE。置信度和类别都用 BCEWithLogitsLoss。关键是正样本筛选遍历每个 ground truth 框找到它中心点落在哪个 grid cell再拿那个 cell 里的 3 个 anchor 去算 IoUIoU 最大的 anchor 负责预测这个框。其余 anchor 里和真实框 IoU 小于 0.5 的位置要被标记为负样本不参与坐标和类别损失只参与置信度损失。这个逻辑如果漏掉模型会倾向于把所有位置都预测成有目标mAP 上不去。4.3 训练骨架一个最小可运行的 train_step 与关键超参以下是一个不依赖第三方检测库的训练循环骨架重点在把两个检测头的 loss 合并起来def train_one_step(model, loss_fn, optimizer, images, targets): det_13, det_26 model(images) loss_13 loss_fn( det_13, targets, anchorsANCHORS_13, stride32, num_classesnum_classes, ) loss_26 loss_fn( det_26, targets, anchorsANCHORS_26, stride16, num_classesnum_classes, ) loss loss_13 loss_26 optimizer.zero_grad() loss.backward() optimizer.step() return loss.item() # 典型训练参数 optimizer torch.optim.SGD(model.parameters(), lr0.001, momentum0.9) scheduler torch.optim.lr_scheduler.MultiStepLR( optimizer, milestones[30, 60], gamma0.1 )逻辑说明两个检测头共享 backbone 的梯度所以loss_13 loss_26直接相加没问题。如果你的数据集里小目标特别多可以给 loss_26 加一个大于 1 的权重常见做法是乘 1.2 到 1.5避免大目标 head 主导整个训练。 参数说明迁移学习场景下初始学习率用0.001比较稳纯随机初始化建议降到0.0005因为随机初始化的梯度噪声大学习率太高容易在第一百步左右炸掉。batch size 在 416 输入下可选 8 或 16显存不够就先冻结 backbone 前几层而不是继续减 batch size。关于训练输出怎么看loss 在前 200 步掉得快是正常的之后进入平台期。如果 loss 一直震荡不降先检查标注文件有没有 0 宽高框再检查 anchor 是否按 mask 分对这两条我几乎每次帮别人排查都能命中一条。5. 训练与部署前的排坑YOLOv3-tiny 更容易踩的五种现场5.1 现象加载完官方权重检测结果全是背景模型能跑通输出也不是零但画出来的框要么只有背景类别要么置信度全部低于 0.1。原因有两类第一类是输入预处理不一致Darknet 训练时把图像缩放到 416x416像素值除以 255而 PyTorch 代码里用了 ImageNet 的 mean/std 归一化直接带偏所有特征分布第二类是权重加载器读到了错误位置最常见的翻转是 BN 的 bias 和 running_mean 读反。解决先去掉 mean/std只做除以 255 的归一化再打印第一层卷积的权重均值和 bias与官方 cfg 对应层的预期值数量级对比。如果差异很大回到加载器检查 BN 参数顺序。5.2 现象训练到第两百步 loss 变 NaNloss 在前期正常突然跳成 NaN后面全部是 NaN。原因集中在三个地方wh 的 exp 出现极端大值、标注里有 0 宽高框、学习率过高。416 输入下如果图像里目标特别小归一化后宽高接近 0MSE 在反传时会产生超大梯度。解决在 loss 里对 wh 的 target 做 clamp限制在[1e-4, 10]范围内数据加载时过滤掉w 0 or h 0的标注随机初始化场景下把学习率降到 0.0005。做这三步之后 NaN 基本不会再出现。5.3 现象小目标一个都检测不到大目标框又偏这是一个有明显规律的现象13x13 检测头好像正常26x26 检测头完全失效。原因基本可以锁定为 anchor mask 分配反了。YOLOv3-tiny 的 13x13 头应该使用六个 anchor 里的后三个大 anchor26x26 头使用前三个小 anchor。如果两套 anchor 调换小目标会被迫用大 anchor 去预测模型学不到合适的偏移。解决对照 cfg 里的 mask 检查 anchor 配置确认 13x13 用的是[81,82,135,169,344,319]26x26 用的是[10,14,23,27,37,58]。调试时可以在 decode 函数里打印每个头激活的 anchor 宽度一眼就能看出有没有分配错。5.4 现象显存占用比预期高一倍416x416 输入、batch size 8 的 YOLOv3-tiny理论显存占用应该在 4GB 到 6GB 之间。如果实测超过 8GB多半不是模型本身的问题。常见原因是训练时所有 backbone 层都保留了梯度且开启了一些额外的中间特征缓存另一个原因是 DataLoader 的num_workers设置过高每个 worker 复制一份数据和模型状态小数据集上更容易放大。解决冻结前几层卷积的参数设置requires_gradFalse把num_workers控制在 4 以内如果开了 AMP检查梯度缩放器是否正常工作half 精度下某些 BN 层会产生额外显存开销。5.5 现象NMS 之后同一个目标出现大量重叠框目标上叠着三四个框置信度都挺高NMS 过滤不掉。原因conf_thresh 设置太低导致背景位置也保留了大量候选框或者类别置信度没有乘上目标置信度而是单独做阈值过滤导致低质量框混进来。解决最终分数用conf * cls_prob相乘后再过滤conf_thresh 在边缘设备场景建议设 0.25精度优先场景设 0.1iou_thresh 设 0.45。这两个阈值不是玄学背后是 precision-recall 的权衡调参时固定一个单变量去扫另一个。6. 进阶把 PyTorch 模型导出 ONNX 并用固定 batch 跑实时推理6.1 用 40 行代码导出静态 ONNX训练完的模型在 PyTorch 里跑实时推理延迟还是偏高。常见做法是导出 ONNX再转成 TensorRT engine。YOLOv3-tiny 结构简单导出非常容易但要注意把 NMS 留在引擎外处理。如果你的部署端支持自定义 NMS 插件可以一起导出如果不支持只导出两个检测头特征图解码和 NMS 在 C 或 Python 侧做兼容性最好。import torch model YOLOv3Tiny(num_classes80) load_darknet_weights(model, yolov3-tiny.cfg, yolov3-tiny.weights) model.eval() dummy torch.randn(1, 3, 416, 416) torch.onnx.export( model, dummy, yolov3-tiny.onnx, opset_version11, input_names[input], output_names[det_13, det_26], dynamic_axesNone, ) print(export done)参数说明dynamic_axesNone意味着导出的是静态 batch、静态分辨率这是 TensorRT 最喜欢的输入形式推理性能最好。如果你的业务要求动态分辨率可以给 h、w 轴设 dynamic但延迟会明显上升边缘设备上通常不值得。导出后可以用 ONNX Runtime 做一次正确性验证把同一张图分别输入 PyTorch 模型和 ONNX 模型对比两个检测头的输出张量最大绝对误差小于 1e-3 就说明导出没丢算子。6.2 部署前必须做的两件事输出对比与精度回落检查部署到边缘设备之前我习惯固定做三个检查第一把 PyTorch 的检测框和 ONNX Runtime 的检测框画在同一张图上肉眼对比位置差异第二在验证集上统计一次 mAP确认导出后精度回落不超过 1 个百分点第三用 TensorRT 的 FP16 模式跑一遍记录延迟和精度变化。YOLOv3-tiny 这种轻量模型FP16 通常能带来约 40% 到 60% 的速度提升精度回落很小。INT8 量化在部分边缘设备上会掉点明显尤其是小目标检测场景我建议先跑 FP16确定精度可接受再考虑 INT8。我自己在导出 ONNX 时吃过一次亏当时为了图省事把 NMS 也写进模型里导出结果部署端不支持那个自定义算子整个推理链路卡了整整一天。后来改成只导出特征图、自己在 C 侧实现 decode 和 NMS问题立刻解决。轻量模型的部署方案越简单越不容易出错前后端各管一段定位问题也更快。希望这篇笔记能帮你把 YOLOv3-tiny 从加载、训练到部署完整走通少踩几次我踩过的坑。本文还有配套的精品资源点击获取
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进