ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

PyTorch实现YOLOv3-tiny:从权重加载到推理部署的完整指南

PyTorch实现YOLOv3-tiny:从权重加载到推理部署的完整指南 简介这份资源是面向深度学习入门者与边缘端部署开发者的 PyTorch 版 YOLOv3-tiny 目标检测工程包聚焦在算力受限环境下实现实时对象检测。压缩包共 22 个文件约 1.17MB以 14 个 Python 脚本为核心辅以 3 张网络结构与锚点聚类示意图、2 个类别名称文件、字体与说明文档覆盖模型定义、数据预处理、训练与推理全链路。代码层面包含预训练、微调与推理脚本yolov3tiny 模块内实现卷积、池化等网络组件及损失函数gather_anchors.py 负责锚点聚类计算build_lmdb.py 用于构建高效 LMDB 数据存储images 目录提供测试样例图片。目前已有 48 人学习。读者可据此快速搭建轻量检测基线理解锚点选取与数据管线设计并在此基础上训练定制化实时检测系统。1. PyTorch 实现 YOLOv3-tiny从权重文件到能跑通的检测工程你手里如果只有一个PyTorch实现YOLOv3-tiny.zip第一反应大概率是解压、找train.py、pip install -r requirements.txt然后被一堆版本冲突和 shape 报错按在地上摩擦。YOLOv3-tiny 这个模型本身不复杂——它只有两个 YOLO 检测头、一个 backbone 加一个轻量 neck参数量在 800 万到 900 万之间权重文件不到 35MBCPU 上单帧推理也能压到几十毫秒。但真正让一线工程师翻车的从来不是网络结构而是环境、权重加载、anchor 匹配和 NMS 后处理这几处“玄学”。这篇内容面向三类人想用 PyTorch 从零复现 YOLOv3-tiny 推理链路的算法工程师、需要把检测模型塞进边缘设备或低算力场景的落地开发者、以及正在做 PyTorch 实战项目、想找一个结构清晰又不至于太重的检测模型练手的学生。我会按“环境搭建 → 网络结构 → 权重加载 → 推理与后处理 → 避坑 → 进阶验证”的顺序把每一步的命令、参数和失败排查讲清楚。你不需要先看完整个 PyTorch 教程但至少要能跑通一个torch.randn的前向传播。2. 环境搭建与 PyTorch 安装别让 CUDA 版本成为第一道坎2.1 用 conda 隔离环境别在 base 里硬装PyTorch 安装教程网上铺天盖地但真正能一次跑通的没几个。核心原则只有一条先确定 CUDA 版本再选 PyTorch 安装命令最后才装其他依赖。如果你用的是 NVIDIA 显卡先在终端执行nvidia-smi看右上角CUDA Version。注意这个版本是驱动支持的最高 CUDA 版本不是你实际要装的 CUDA Toolkit 版本。PyTorch 官方 wheel 包自带 CUDA runtime你不需要单独装 CUDA Toolkit只需要驱动版本足够高。我一般用 conda 建一个独立环境Python 版本选 3.9 或 3.10这两个版本对 PyTorch 和 OpenCV 的兼容性最稳。命令如下conda create -n yolov3tiny python3.10 -y conda activate yolov3tiny创建完环境后去 PyTorch 官网的 previous versions 页面找对应 CUDA 版本的安装命令。假设你的驱动支持 CUDA 11.8命令大概是pip install torch2.1.0 torchvision0.16.0 --index-url https://download.pytorch.org/whl/cu118如果你没有 NVIDIA 显卡或者只是想在 CPU 上先跑通推理直接装 CPU 版pip install torch2.1.0 torchvision0.16.0 --index-url https://download.pytorch.org/whl/cpu装完后用下面这段代码验证import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else CPU mode)如果torch.cuda.is_available()返回False先别急着重装。常见原因有三个驱动版本太低、装成了 CPU 版、conda 环境里混入了其他 torch 包。用pip list | grep torch确认一下如果看到torch和torchvision版本不匹配直接pip uninstall后重装。提示WSL 环境下装 PyTorch 时nvidia-smi能显示显卡不代表 WSL 内 CUDA 可用。需要在 Windows 侧装好支持 WSL 的驱动然后在 WSL 里执行nvidia-smi验证。如果 WSL 里看不到显卡检查 Windows 驱动版本是否满足 WSL CUDA 要求。2.2 依赖清单与版本锁定YOLOv3-tiny 的推理依赖不多但版本冲突往往出在numpy和opencv-python上。我一般会固定这几个版本依赖包推荐版本作用torch2.1.0模型推理与权重加载torchvision0.16.0图像预处理与 NMSopencv-python4.8.1.78图像读取、绘制框numpy1.24.4数组运算pillow10.0.1图像格式转换tqdm4.66.1推理进度显示安装命令pip install opencv-python4.8.1.78 numpy1.24.4 pillow10.0.1 tqdm4.66.1这里有个血泪经验opencv-python4.9 以上版本在某些 Linux 发行版上会依赖libGL如果你在服务器上跑可能报ImportError: libGL.so.1: cannot open shared object file。解决办法是装opencv-python-headless或者用apt install libgl1补上系统库。我一般直接换 headless 版本少一个系统依赖少一个坑。2.3 验证环境是否真的能跑推理环境装好后别急着下载权重。先用一个随机张量测一下前向传播是否正常import torch import torch.nn as nn # 模拟一个卷积层验证 CUDA 是否真的能跑 device torch.device(cuda if torch.cuda.is_available() else cpu) conv nn.Conv2d(3, 16, kernel_size3, padding1).to(device) x torch.randn(1, 3, 416, 416).to(device) y conv(x) print(y.shape) # 应该是 torch.Size([1, 16, 416, 416])如果这一步报CUDA out of memory说明显卡显存被其他进程占了用nvidia-smi看一下必要时kill掉无关进程。如果报RuntimeError: CUDA error: no kernel image is available for execution on the device说明 PyTorch 版本和显卡算力不匹配需要换对应 CUDA 版本的 wheel 包。3. YOLOv3-tiny 网络结构拆解为什么它能在低算力设备上跑3.1 backbone 与 neck 的层数分配YOLOv3-tiny 的结构比 YOLOv3 简单得多。它没有残差块没有 CSP 结构整个 backbone 就是 7 个卷积层加 5 个最大池化层最后接两个 YOLO 检测头。具体来说前 6 个卷积层负责浅层特征提取通道数从 3 逐步升到 256第 7 个卷积层后接一个最大池化然后分成两路一路继续下采样接第 8 到第 10 个卷积层输出 13×13 的特征图接第一个 YOLO 头另一路从第 7 层后上采样与第 8 层的特征拼接再接第 11 到第 13 个卷积层输出 26×26 的特征图接第二个 YOLO 头这种“一大一小”双检测头设计让 YOLOv3-tiny 能同时检测大目标和小目标。13×13 的特征图感受野大适合检测大物体26×26 的特征图分辨率高适合检测小物体。但它的短板也很明显小目标检测精度远不如 YOLOv3因为浅层特征没有足够的语义信息。用 PyTorch 定义这个结构时我一般用一个nn.ModuleList管理卷积层用nn.MaxPool2d和nn.Upsample处理下采样和上采样。核心代码如下import torch import torch.nn as nn class YOLOv3Tiny(nn.Module): def __init__(self, num_classes80): super().__init__() self.num_classes num_classes # backbone: 7 个卷积 5 个池化 self.backbone nn.Sequential( nn.Conv2d(3, 16, 3, 1, 1, biasFalse), nn.BatchNorm2d(16), nn.LeakyReLU(0.1, inplaceTrue), nn.MaxPool2d(2, 2), # ... 中间层省略按实际结构补全 ) # 检测头 1: 13x13 self.head1 nn.Conv2d(256, 3 * (num_classes 5), 1, 1, 0) # 检测头 2: 26x26 self.head2 nn.Conv2d(128, 3 * (num_classes 5), 1, 1, 0) def forward(self, x): # 前向传播逻辑 pass这里的关键参数是num_classes和每个检测头的输出通道数。YOLOv3-tiny 每个 grid cell 预测 3 个 anchor每个 anchor 输出(num_classes 5)个值其中 5 是x, y, w, h, confidence。所以检测头的输出通道数必须是3 * (num_classes 5)。如果你用的是 COCO 数据集num_classes80输出通道就是3 * 85 255。3.2 anchor 机制与输出解码YOLOv3-tiny 的 anchor 是预先聚类好的。COCO 数据集上常用的 6 个 anchor 是(10,14), (23,27), (37,58), (81,82), (135,169), (344,319)前 3 个给 26×26 的检测头后 3 个给 13×13 的检测头。模型输出的tx, ty, tw, th是相对值需要解码成绝对坐标def decode_boxes(pred, anchors, stride): # pred: [batch, 3, H, W, num_classes 5] # anchors: 当前检测头对应的 3 个 anchor # stride: 13 或 26 batch, _, H, W, _ pred.shape grid_x torch.arange(W).view(1, 1, W, 1).float() grid_y torch.arange(H).view(1, H, 1, 1).float() anchor_w torch.tensor(anchors)[:, 0].view(1, 3, 1, 1).float() anchor_h torch.tensor(anchors)[:, 1].view(1, 3, 1, 1).float() # sigmoid 激活中心点偏移 pred[..., 0] (torch.sigmoid(pred[..., 0]) grid_x) * stride pred[..., 1] (torch.sigmoid(pred[..., 1]) grid_y) * stride # 宽高用 exp 激活 pred[..., 2] torch.exp(pred[..., 2]) * anchor_w pred[..., 3] torch.exp(pred[..., 3]) * anchor_h return pred这段代码里最容易出错的是grid_x和grid_y的维度。如果维度对不上广播机制会静默产生错误结果而不是报错。我一般会在解码后打印几个框的坐标确认数值在合理范围内比如 x 在 0 到 416 之间。如果出现负数或超过图像尺寸的值大概率是 grid 维度搞错了。3.3 用 torchsummary 检查参数量和层结构在正式加载权重前建议先用torchsummary看一下模型结构是否符合预期from torchsummary import summary model YOLOv3Tiny(num_classes80) summary(model, input_size(3, 416, 416))输出会列出每一层的输出 shape 和参数量。YOLOv3-tiny 的总参数量应该在 8.7M 左右如果差太多说明你的结构定义有问题。常见错误包括漏了某个卷积层、通道数写错、上采样方式不对。这一步花 5 分钟能省掉后面几小时的调试时间。4. 权重加载与推理从 .weights 到 PyTorch state_dict4.1 Darknet 权重格式解析YOLOv3-tiny 的官方权重是 Darknet 格式的.weights文件不是 PyTorch 的.pth。这个文件的结构是前 5 个 int32 是头部信息主版本、次版本、修订版本、训练图像尺寸、是否见过真值后面按层顺序存储卷积层的权重和 BN 参数。解析代码如下import numpy as np import torch def load_darknet_weights(model, weights_path): with open(weights_path, rb) as f: # 前 5 个 int32 是头部 header np.fromfile(f, dtypenp.int32, count5) weights np.fromfile(f, dtypenp.float32) ptr 0 for m in model.modules(): if isinstance(m, torch.nn.Conv2d): # 卷积层权重 conv_w torch.from_numpy( weights[ptr:ptr m.weight.numel()] ).view_as(m.weight) m.weight.data.copy_(conv_w) ptr m.weight.numel() # 如果有 BN接着读 bias、running_mean、running_var、weight if m.bias is not None: bias torch.from_numpy( weights[ptr:ptr m.bias.numel()] ).view_as(m.bias) m.bias.data.copy_(bias) ptr m.bias.numel() elif isinstance(m, torch.nn.BatchNorm2d): # BN 的 4 个参数 for param in [m.bias, m.weight, m.running_mean, m.running_var]: param.data.copy_( torch.from_numpy( weights[ptr:ptr param.numel()] ).view_as(param) ) ptr param.numel() print(fLoaded {ptr} parameters from {weights_path})这段代码的核心逻辑是按模型定义的顺序依次从weights数组里读取对应数量的浮点数填充到每个层的参数里。最容易翻车的地方是BN 参数的顺序。Darknet 里 BN 的存储顺序是bias, weight, running_mean, running_var而 PyTorch 的state_dict里顺序可能不同。如果你加载后推理结果全是乱框大概率是 BN 参数顺序搞反了。注意加载完权重后一定要调用model.eval()并把 BN 层设为评估模式。否则 BN 会继续用当前 batch 的统计量导致推理结果不稳定。4.2 图像预处理与推理命令推理前的预处理包括letterbox 缩放、归一化、通道转换。letterbox 的作用是保持图像宽高比用灰色填充到 416×416。代码如下import cv2 import numpy as np def letterbox(img, new_shape416, color(128, 128, 128)): shape img.shape[:2] r min(new_shape / shape[0], new_shape / shape[1]) new_unpad (int(round(shape[1] * r)), int(round(shape[0] * r))) dw, dh new_shape - new_unpad[0], new_shape - new_unpad[1] dw / 2 dh / 2 img cv2.resize(img, new_unpad, interpolationcv2.INTER_LINEAR) top, bottom int(round(dh - 0.1)), int(round(dh 0.1)) left, right int(round(dw - 0.1)), int(round(dw 0.1)) img cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, valuecolor) return img, r, (left, top) # 使用示例 img cv2.imread(test.jpg) img_letterbox, ratio, pad letterbox(img, 416) img_rgb cv2.cvtColor(img_letterbox, cv2.COLOR_BGR2RGB) img_tensor torch.from_numpy(img_rgb).float().permute(2, 0, 1) / 255.0 img_tensor img_tensor.unsqueeze(0)这里有几个参数需要留意new_shape416是 YOLOv3-tiny 的标准输入尺寸但如果你显存不够可以改成 320 或 288精度会降但速度会升。color(128,128,128)是填充色Darknet 默认用 0.5 灰度对应 128。ratio和pad在后续把框映射回原图时要用到别丢掉。推理命令很简单model.eval() with torch.no_grad(): pred model(img_tensor)如果显存够可以把img_tensor放到 GPU 上同时model.to(cuda)。推理时间在 RTX 3060 上大约是 3 到 5 毫秒在 CPU 上大约是 30 到 50 毫秒。4.3 NMS 后处理与框映射模型输出的是原始预测值需要经过解码、置信度过滤、NMS 才能得到最终框。NMS 我一般直接用torchvision.ops.nms比自己写更稳from torchvision.ops import nms def postprocess(pred, conf_thres0.5, iou_thres0.4): # pred: [1, 2535, 85] 或类似形状 pred pred[pred[..., 4] conf_thres] if pred.shape[0] 0: return None # 类别置信度 class_conf, class_pred pred[..., 5:].max(1, keepdimTrue) conf pred[..., 4:5] * class_conf # 框坐标转换 boxes pred[..., :4] # NMS keep nms(boxes, conf.squeeze(), iou_thres) return boxes[keep], conf[keep], class_pred[keep]conf_thres0.5和iou_thres0.4是两个需要根据场景调的参数。如果你发现漏检多把conf_thres降到 0.3如果误检多升到 0.6。iou_thres控制重叠框的合并程度密集场景可以降到 0.3稀疏场景可以升到 0.5。框映射回原图的公式是# boxes 是相对 416x416 的坐标 boxes[:, [0, 2]] (boxes[:, [0, 2]] - pad[0]) / ratio boxes[:, [1, 3]] (boxes[:, [1, 3]] - pad[1]) / ratio这一步如果忘了减pad框会整体偏移。我见过有人调了半天模型最后发现是 letterbox 的 padding 没减掉。5. 避坑与排查YOLOv3-tiny 落地时最容易翻车的 5 个点5.1 现象推理结果全是乱框置信度极低原因权重加载时 BN 参数顺序错误或者model.eval()没调用。Darknet 的 BN 存储顺序和 PyTorch 的state_dict顺序不一致如果按 PyTorch 的顺序读running_mean和running_var会错位。解决在加载权重前打印模型每一层的参数名和 shape对照 Darknet 的层顺序逐个核对。加载完后用一张已知图片测试如果置信度普遍低于 0.1基本可以确定是权重加载问题。另外确保model.eval()在推理前调用且torch.no_grad()包裹推理过程。5.2 现象CUDA out of memory但显存明明够原因PyTorch 默认会缓存显存如果之前跑过其他模型缓存没释放。或者img_tensor没有detach()计算图一直挂着。解决在推理循环里加torch.cuda.empty_cache()或者把img_tensor用with torch.no_grad()包住。如果还是不够把输入尺寸从 416 降到 320显存占用能减少约 40%。另外检查是不是model和img_tensor不在同一个设备上跨设备传输会额外占显存。5.3 现象小目标检测效果差几乎全漏原因YOLOv3-tiny 的 26×26 检测头虽然分辨率高但浅层特征语义信息弱。如果输入尺寸降到 320 以下小目标特征会进一步丢失。解决保持输入尺寸在 416 或更高同时把conf_thres降到 0.3。如果还是不行考虑在 52×52 的特征图上加第三个检测头但这会改变网络结构需要重新训练。另一个办法是用 TTA测试时增强把图片放大 1.2 倍再推理一次合并结果。5.4 现象NMS 后框大量重叠同一个目标出多个框原因iou_thres设得太高或者类别置信度和 objectness 的乘积计算错误。解决把iou_thres从 0.5 降到 0.3 到 0.4。检查conf pred[..., 4:5] * class_conf这一步确保是乘法而不是加法。如果同一个目标被分成不同类别说明分类头输出有问题检查num_classes是否和权重文件匹配。5.5 现象CPU 推理速度远低于预期原因PyTorch 默认用单线程推理没有开 MKL 或 OpenMP 加速。或者模型没有转成torch.jit脚本。解决设置torch.set_num_threads(4)或更高具体取决于 CPU 核心数。如果追求极致速度用torch.jit.trace把模型转成 TorchScript推理速度能提升 20% 到 30%。另外把model.eval()和torch.no_grad()都加上避免不必要的计算。6. 进阶验证用 ONNX 导出和量化把 YOLOv3-tiny 压到极致6.1 导出 ONNX 并验证数值一致性PyTorch 转 ONNX 是部署到 TensorRT、OpenVINO 或 NCNN 的必经之路。导出命令如下import torch.onnx model.eval() dummy_input torch.randn(1, 3, 416, 416) torch.onnx.export( model, dummy_input, yolov3_tiny.onnx, input_names[input], output_names[output], opset_version11, dynamic_axes{input: {0: batch}, output: {0: batch}} )opset_version11是兼容性最好的版本支持大多数推理引擎。导出后用onnxruntime验证数值一致性import onnxruntime as ort import numpy as np sess ort.InferenceSession(yolov3_tiny.onnx) onnx_out sess.run(None, {input: dummy_input.numpy()}) torch_out model(dummy_input).detach().numpy() print(np.max(np.abs(onnx_out[0] - torch_out)))如果差值大于 1e-4说明导出过程中有算子不兼容。常见问题是nn.Upsample在 opset 11 里需要用modenearest如果用了bilinear可能会报错。6.2 动态量化与推理速度对比PyTorch 支持动态量化把卷积层的权重从 float32 转成 int8模型体积能压到原来的 1/4CPU 推理速度提升 2 到 3 倍。命令如下import torch.quantization model.eval() quantized_model torch.quantization.quantize_dynamic( model, {torch.nn.Conv2d, torch.nn.Linear}, dtypetorch.qint8 ) torch.save(quantized_model.state_dict(), yolov3_tiny_quantized.pth)量化后精度会掉 1 到 3 个 mAP 点具体取决于数据集。如果对精度要求高可以用量化感知训练QAT在训练时模拟量化误差精度损失能控制在 1 个点以内。6.3 一个我常用的验证习惯每次改完模型结构或后处理逻辑我都会用同一张测试图跑一遍把框画出来存成图片然后和上一版结果对比。如果框的位置偏移超过 5 个像素或者置信度变化超过 0.1就说明改动引入了问题。这个习惯帮我省了很多“改完不知道对不对”的时间。另外我会在代码里固定随机种子确保每次推理的输入完全一致torch.manual_seed(42) np.random.seed(42)YOLOv3-tiny 这个模型不大但要把推理链路跑通、跑稳需要关注的细节不少。从环境搭建到权重加载从 anchor 解码到 NMS每一步都有坑。我一般会先把推理跑通再考虑训练和部署。如果你也在做类似的项目建议先把单张图片的推理结果调对再批量跑视频或摄像头。希望帮到你。本文还有配套的精品资源点击获取
返回列表