ARTICLE DETAIL

资讯详情

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

浏览器端YOLOv5实时检测:PyTorch到TF.js的完整转换指南

浏览器端YOLOv5实时检测:PyTorch到TF.js的完整转换指南 简介YOLOv5与TensorFlow.js的整合示例包面向希望在Web端实现实时目标检测的开发者解决模型部署到浏览器或Node.js环境时的跨平台集成问题。资源共29个文件整体仅88KB结构精炼前端以HTML、CSS、JS页面资源为主后端采用Python脚本辅助sh安装部署脚本及cfg、yaml配置文件方便在不同框架下快速迁移。目前已有289人学习关注适合具备一定YOLO基础、希望延伸至Web端推理或轻量级部署的开发者。压缩包内提供了基于FastAPI、Flask、Bottle等多种后端框架的接入示例并包含HLS流式处理相关前端静态资源同时整理了gitignore、license、CodeCheck等工程规范文件便于项目维护和合规使用。借助这套资源可以减少环境配置和前后端联调中的重复劳动快速搭建可运行的浏览器端YOLOv5检测应用。 最近有个需求反复被问到训练好的YOLOv5模型能不能不部署到服务器直接在浏览器里跑实时检测一开始我觉得这需求有点奇怪毕竟常规路线都是走API调用但实际落地场景其实挺多的——比如内部工具链里的摄像头质检演示、纯前端的数据标注辅助、甚至一些需要隐私保护的本地化识别场景浏览器端推理确实能省掉一大截服务器带宽和中间链路。所以就有了这个YOLOv5-RT-TFJS项目把PyTorch训练好的YOLOv5权重转换到TensorFlow.js格式在浏览器里用WebGL后端做实时目标检测。这篇文章就把整个链路的完整步骤、转换踩坑、前后处理细节和实测性能一次说清楚。1. 为什么非要折腾到浏览器里跑场景驱动与技术选型1.1 省掉服务器的诱惑与代价浏览器端跑推理最大的好处有两点零部署成本、零传输延迟。模型权重分片加载到本地后图像数据不需要上传到服务器摄像头画面直接在浏览器内存里完成推理整个耗时只有模型推理时间加一点点前后处理时间没有任何网络抖动。对Demo演示、快速原型验证这类场景来说体验非常好。代价也很明显浏览器能用的算力是受限的。WebGL后端虽然能用GPU但抽象层比CUDA厚重运行效率跟原生PyTorch比有差距。所以这个方案适合轻量级模型比如YOLOv5s、YOLOv5n级别跑YOLOv5m都开始吃力。如果目标是YOLOv5l甚至x那我建议死了这条心老老实实上后端服务。1.2 为什么选TensorFlow.js而不是ONNX WebAssembly有些同学会问PyTorch模型不是可以直接装onnxruntime-web跑吗这条路也通但实际体验下来有几个问题onnxruntime-web对WebGL的算子覆盖不如TensorFlow.js全面YOLOv5导出ONNX后有些自定义算子需要额外处理。TensorFlow.js提供了tf.browser.fromPixels这类专门优化过的浏览器API视频帧转Tensor的效率更高。TFJS的executeAsync支持多输出一次性拉取YOLOv5的三个检测头输出刚好能一起拿回来减少JS-WebGL边界的数据拷贝。所以我选择的核心链路是PyTorch权重 → ONNX → TensorFlow SavedModel → TF.js这条链路在github上已经有很多成功案例踩坑成本相对低。1.3 项目最终呈现的效果输入可以是video标签接摄像头也可以是图片上传。视频流模式用requestAnimationFrame驱动抽帧每帧resize到640x640推理得到三个尺度的特征图前端解码加NMS最后用canvas把检测框画出来。选配项包括置信度阈值、IOU阈值、类别名导入。整个项目依赖全部在npm里打包后可以丢到任意静态服务器。2. 模型转换链路从PyTorch权重到TF.js的全流程2.1 转换环境准备转换过程最好在Linux或macOS环境下做Windows也能跑但onnx转tf时偶尔会遇到路径编码问题。建议单独建一个Python虚拟环境避免污染已有的深度学习环境python -m venv yolov5_tfjs_env source yolov5_tfjs_env/bin/activate pip install torch torchvision onnx onnx2tf tensorflow tensorflowjs注意版本坑tensorflow和tensorflowjs的版本要匹配。我这里用的是tensorflow2.13.1搭配tensorflowjs4.15.0实测稳定。onxx2tf默认走TF 2.x的SavedModel导出版本太新反而可能出现算子兼容性问题。2.2 第一步PyTorch转ONNX如果你用的是ultralytics官方YOLOv5仓库导出命令很简单git clone https://github.com/ultralytics/yolov5 cd yolov5 python export.py --weights yolov5s.pt --include onnx --img 640这里关键是--img 640它决定了模型输入分辨率。浏览器端处理时输入Tensor必须是这个尺寸。如果你的使用场景分辨率不固定想要动态尺寸可以加--dynamic参数但TFJS的WebGL后端对动态shape支持不太好强烈建议固定尺寸导出。我一般固定640在浏览器端做letterbox处理。导出完成后用onnx.checker验证一下模型完整性import onnx model onnx.load(yolov5s.onnx) onnx.checker.check_model(model) print(ONNX模型检查通过)2.3 第二步ONNX转TensorFlow SavedModel这里推荐onnx2tf这个工具它比老牌的onnx-tf更新对ONNX算子覆盖更全而且会自动处理一些TF不支持的算子onnx2tf -i yolov5s.onnx -o yolov5s_savedmodel执行完会在yolov5s_savedmodel目录下生成saved_model.pb和variables目录。这里有个小细节onnx2tf默认会把输出节点的名字改成类似Identity、Identity_1这样的名称后面在TFJS里executeAsync取输出时要靠索引而不是名字所以不用太在意输出节点名称。2.4 第三步SavedModel转TF.js格式最后一步用tensorflowjs官方转换器tensorflowjs_converter \ --input_formattf_saved_model \ --output_formattfjs_graph_model \ --signature_nameserving_default \ yolov5s_savedmodel \ yolov5s_tfjs转换完成后yolov5s_tfjs目录下会出现model.json和一组group1-shard1ofN.bin权重分片文件。这里的坑是默认权重分片可能比较大浏览器加载时如果服务器不支持Range请求分片加载会失败。一般在Web服务器配置里确认支持Accept-Ranges: bytes即可或者用更小的分片重新导出tensorflowjs_converter ... --weight_shard_size_bytes 4194304这样每个分片4MB加载更稳。2.5 实测踩坑OpSet版本与算子兼容我在转换过程中遇到最烦的问题就是ONNX的OpSet版本。YOLOv5官方导出默认用的是OpSet 12或17如果onnx2tf对应的TF版本解析不了某些算子尤其是Resize和Mul的组合会直接报错。解决方式是导出ONNX时手动指定OpSetpython export.py --weights yolov5s.pt --include onnx --opset 12实测OpSet 12搭配onnx2tf成功概率最高。另外如果模型训练时替换了激活函数比如用了SiLU的变体转换时也可能出问题这种情况建议先转成标准的ReLU激活再导出虽然会损失一点精度但转换通畅。3. 浏览器端推理管线搭建加载TF.js模型与视频流处理3.1 TF.js运行时配置前端项目直接用npm安装依赖npm install tensorflow/tfjs tensorflow/tfjs-backend-webgl初始化时强制指定WebGL后端import * as tf from tensorflow/tfjs; import tensorflow/tfjs-backend-webgl; await tf.setBackend(webgl); await tf.ready();这里有个容易犯的错误如果同时引入了tensorflow/tfjs-node或者本地后端的包初始化时可能自动选择CPU后端导致推理速度慢好几倍。调试时可以通过tf.getBackend()确认当前后端。3.2 加载图模型与首次预热图模型的加载方式和LayersModel不同必须用tf.loadGraphModelconst model await tf.loadGraphModel(/models/yolov5s_tfjs/model.json);加载完毕后建议做一次预热推理因为WebGL着色器是首次用到时才编译预热能避免第一帧卡顿const dummy tf.zeros([1, 640, 640, 3], int32); await model.executeAsync(dummy); dummy.dispose();注意这里用int32类型模拟输入像素值。YOLOv5导出的模型如果NCHW还是NHWC布局写死在图里输入类型也会不同如果你发现模型报错改成tf.zeros([1, 3, 640, 640], float32)试一下。我们导出时通过onnx2tf的--output_signaturedefs参数强制NHWC输出让浏览器端处理更顺畅。3.3 视频帧捕获与预处理用原生navigator.mediaDevices.getUserMedia拿摄像头流然后把video元素挂到页面上隐藏起来每帧用requestAnimationFrame抽取。核心预处理函数function preprocess(video, targetSize 640) { const tensor tf.browser.fromPixels(video) .resizeBilinear([targetSize, targetSize]) .expandDims(0) .div(255.0); return tensor; }tf.browser.fromPixels是把HTMLVideoElement直接转成Tensor的内置API内部走的是WebGL纹理拷贝效率远高于先画到canvas再getImageData。resizeBilinear这一步可以换成resizeNearestNeighbor实测对检测精度影响不大但速度略快看你的场景取舍。3.4 executeAsync与输出张量管理YOLOv5导出后有三个输出对应三个检测头。使用executeAsync一次性取回全部输出const input preprocess(video); const outputs await model.executeAsync(input); // outputs是数组 // outputs[0]: shape [1, 20, 20, 255] // outputs[1]: shape [1, 40, 40, 255] // outputs[2]: shape [1, 80, 80, 255]以COCO数据集为例类别数80每个网格位置预测3个anchor所以通道数是(5 80) * 3 255。如果你的模型类别数是别的值比如只检测几个类别这个数字会变。判断通道数的公式别记错(5 numClasses) * numAnchors。executeAsync返回的Tensor用完必须手动dispose()否则会内存泄漏。在每帧处理函数末尾统一清理input.dispose(); outputs.forEach(t t.dispose());4. 后处理细节anchors、解码公式与NMS实战4.1 为什么导出结构里没有NMS很多人在转换前会问YOLOv5官方导出时不是有--include nms参数吗加上这个确实会生成端到端带NMS的模型输出直接是检测结果。但实测在TFJS里跑这种模型有两个问题一是图里的NMS算子是用TF的自定义操作实现的WebGL后端支持不稳定二是在浏览器端做NMS其实很灵活阈值调起来方便。所以我选择导出纯净的raw模型把NMS留在JavaScript里做。4.2 解码公式与anchors配置YOLOv5使用的anchor是预设在模型里的与训练配置一致。YOLOv5s在640分辨率下三类anchorconst ANCHORS [ [[10, 13], [16, 30], [33, 23]], // 对应80x80特征图 [[30, 61], [62, 45], [59, 119]], // 对应40x40特征图 [[116, 90], [156, 198], [373, 326]] // 对应20x20特征图 ]; const STRIDES [8, 16, 32];特征图尺寸和stride是成反比的输入640stride 8对应80x80stride 16对应40x40stride 32对应20x20。这个匹配关系就是anchor配比的核心逻辑。解码的核心公式每个网格位置、每个anchorconst cellX (Math.sigmoid(xy[0]) * 2 - 0.5 gridX) * stride; const cellY (Math.sigmoid(xy[1]) * 2 - 0.5 gridY) * stride; const boxW Math.pow(Math.sigmoid(wh[0]) * 2, 2) * anchorW; const boxH Math.pow(Math.sigmoid(wh[1]) * 2, 2) * anchorH; const objScore Math.sigmoid(obj); const classScores classLogits.map(logit Math.sigmoid(logit) * objScore);公式里的乘2减0.5和乘2平方是YOLOv5在v6.0之后引入的解码修正让中心点偏移范围从0到1变成-0.5到1.5能提高收敛速度。如果碰到老版本模型v5.0之前的权重解码公式不同中心点是sigmoid加网格偏移宽高是exp乘以anchor使用前要确认你的权重版本。4.3 Tensor操作还是循环操作这里有一个性能分水岭直接在Tensor层面做解码利用WebGL并行计算速度很快但代码复杂用数组循环做解码代码直观但纯JS单线程跑起来很慢。实测在60fps的摄像头流上每帧有(20*20 40*40 80*80) * 3 25200个候选框纯JS循环解码每帧耗时约30ms基本就要丢帧了。我的做法是先用Tensor操作做一次简化筛选只保留置信度高于阈值的网格索引再对这些索引用普通JS数组做精确解码。这样能把进入循环的候选框数量从25200降到几百个解码耗时控制在2ms以内。具体流程// 1. 先取置信度最大值和对应的类别索引降低维度 const [maxScores, maxIndices] tf.topk(classLogits, 1); // 2. 再和objScore相乘过滤低于阈值的 const mask tf.greater(maxScores, minConf); const validIndices tf.where(mask).dataSync(); // 3. 后面只对validIndices对应的网格做JS解码4.4 NMS实现要点与自适应滤波NMS我用的是tf.image.nonMaxSuppression封装好的TensorFlow.js函数没必要自己造轮子const selected await tf.image.nonMaxSuppression( boxes, // shape [numBoxes, 4] scores, // shape [numBoxes] maxOutputSize, // 最大输出框数 iouThreshold, scoreThreshold );但这里有个大坑nonMaxSuppression在WebGL后端底层是走CPU实现的如果输入候选框超过几千个反而会出现明显卡顿。所以一定要在NMS之前先用minConf阈值过滤掉低置信度的框。这个过滤逻辑和前一步的topk筛选是配合使用的。NMS的IOU阈值我的经验值是0.45到0.5之间。目标重叠多就选0.4目标稀疏就选0.5。置信度阈值则看场景演示场景我习惯0.25因为要容忍更多误检来展示识别能力实际工程落地建议0.4起步尤其是误检代价高的场景。4.5 从模型坐标到画布坐标解码得到的框坐标是相对于640x640输入图的。如果实际画面不是正方形我处理的方法是先按照宽高比把原始视频帧缩放到640x640的letterbox尺寸四周用灰边填充解码后的坐标需要做一次反向映射去掉灰边还原到原始视频坐标。function letterboxReverse(box, originalW, originalH, targetSize 640) { const scale Math.min(targetSize / originalW, targetSize / originalH); const resizedW originalW * scale; const resizedH originalH * scale; const offsetX (targetSize - resizedW) / 2; const offsetY (targetSize - resizedH) / 2; return [ (box[0] - offsetX) / scale, (box[1] - offsetY) / scale, (box[2] - offsetX) / scale, (box[3] - offsetY) / scale ]; }如果直接用resizeBilinear把非正方形画面拉伸到640x640虽然省事但检测框的宽高比会变形后面画框时还要再做一次等比缩放反而更容易出问题。5. 实时性能优化与踩坑记录5.1 前端实测数据GPU切换与帧率表现在我测试的几台设备上使用WebGL后端的效果如下设备GPU推理耗时ms帧率FPS备注MacBook Pro M1Apple M1 GPU18-25约30稳定偶尔波动ThinkPad X1 CarbonIntel Iris Xe30-40约20发热后掉帧严重台式机RTX 3070NVIDIA CUDA28-35约25浏览器端吃显卡驱动优化安卓旗舰手机Adreno 66035-50约15作为参考注意RTX 3070的推理耗时没有比核显快多少原因是WebGL后端在NVIDIA驱动上走的是GLSL着色器不是WebGPU路线计算密度上不去。想追求极致性能可以关注WebGPU后端但目前还不太成熟。5.2 每帧推理的pipeline化浏览器端实时推理最怕的是主线程被阻塞。我调整代码结构时把几个操作做成了同步链async function detectFrame() { if (isProcessing) return; // 防止上一帧还没结束就处理下一帧 isProcessing true; const input preprocess(video); const outputs await model.executeAsync(input); const boxes decodeOutputs(outputs); const filtered nmsFilter(boxes); drawBoxes(filtered); input.dispose(); outputs.forEach(t t.dispose()); isProcessing false; } function loop() { requestAnimationFrame(loop); detectFrame(); }isProcessing标志位是关键否则每次requestAnimationFrame都会触发新一轮推理GPU负载会线性增长帧率反而崩溃。实际测试中每帧间隔保持在50ms左右也就是20FPS展示效果完全够用还能留出余量给其他前端逻辑。5.3 稳定性专项内存泄漏排查浏览器端跑推理最阴间的bug就是内存泄漏。跑了几分钟后页面卡死Chrome的Performance面板显示GPU内存持续上升。排查后发现两个问题一是tf.browser.fromPixels返回的Tensor如果不及时dispose()WebGL纹理内存不会立即释放。我在每帧末尾统一处理所有Tensor后内存曲线平稳了。二是executeAsync输出的Tensor数组里有些中间Tensor例如topk产生的索引Tensor虽然不在最终输出里但如果没有用tf.tidy()包裹同样会泄漏。我的做法是外层包一个tf.tidy()让所有中间Tensor自动回收const results await tf.tidy(() { const input preprocess(video); return model.executeAsync(input); });还有个小技巧如果使用了dataSync()比如NMS前要取置信度数组这个方法会阻塞主线程并把数据从GPU拷贝回CPU如果每一帧都调用就会明显卡顿。优化方案是只在置信度阈值过滤的候选框数量变化时才更新数据或者干脆用await tensor.data()这个异步方案释放主线程。5.4 模型尺寸与加载速度的取舍yolov5s.pt大约28MB导出成TFJS后权重分片总大小大约在15MB左右网络加载速度影响明显。如果你的场景对内网加载速度敏感有两个优化方向一是改用yolov5n模型TFJS权重只有6-8MB加载快推理速度也快代价是mAP从s版的37.4掉到28.0左右COCO数据集小目标检测能力下降明显。适合摄像头距离近、目标大的场景。二是打开gzip压缩。TFJS权重分片是二进制格式压缩率很高实测gzip后能压缩到原来的一半左右。确保nginx或任意静态服务器对.bin和.json启用了gzip_static或brotli。5.5 兼容性自查与常见白屏问题浏览器端项目最怕的就是白屏而且往往控制台都没报错。我遇到的几个典型情况模型加载时fetch失败检查模型路径是否被构建工具Vite/Webpack正确处理。我在Vite项目里直接把模型目录放到了public目录而不是通过import引入。WebGL上下文丢失。长页面里如果canvas标签过多GPU上下文会被浏览器回收。给canvas单独设置尺寸而不是继承默认并且避免在隐藏页面里继续跑推理。跨域问题不要把模型放在CDN上又用不同域名访问TFJS的fetch不支持跨域就没法加载。最好和页面同源。6. 一个基于经验的自检清单与扩展思路最后分享一个我每次调试这类项目都会过的自检清单能省掉大量定位问题的时间模型输入尺寸检查model.json里的input shape是不是[1, 640, 640, 3]如果变成[1, 3, 640, 640]说明转换时没有NHWC转换成功。输出通道数三个head的输出最后一个维度是不是(5类别数)*3如果对不上大概率是导出时类别数与训练时不一致。置信度阈值过滤在NMS前必须过滤否则候选框数量太多CPU端NMS会成为瓶颈。内存曲线跑10分钟看Chrome的内存曲线是不是平稳的有任何单调上涨都要怀疑Tensor没释放。这个项目的后续扩展思路其实很多。一个是把模型切分成不同尺寸的版本根据用户设备的WebGL能力动态选择加载对应版本另一个是尝试WebGPU后端推理速度有翻倍空间还有一个是把标注结果直接回传训练集形成浏览器端标注-训练-部署的闭环。我自己下一步打算测试一下加了TinyTrack结构蒸馏出的YOLOv5n在实际浏览器端的表现到时候再开一篇单独聊。本文还有配套的精品资源点击获取
返回列表