DAMO-YOLO TinyNAS代码实例:EagleEye自定义类别训练与ONNX导出完整流程

DAMO-YOLO TinyNAS代码实例:EagleEye自定义类别训练与ONNX导出完整流程

1. 为什么选EagleEye?轻量、快、稳、全本地

你有没有遇到过这样的问题:想在边缘设备或普通工作站上跑一个目标检测模型,但YOLOv8太重,YOLOv5精度又不够,NanoDet推理快却泛化弱,而部署时还要反复折腾ONNX、TensorRT、OpenVINO……最后卡在“能跑”和“跑得好”之间?

EagleEye不是另一个“又一个YOLO变体”。它是达摩院DAMO-YOLO TinyNAS架构落地的实战组合——不靠堆显存,不靠升分辨率,而是用神经架构搜索(TinyNAS)从头筛出最适合低延迟场景的子网络结构。它在单张RTX 4090上实测平均推理耗时17.3ms(输入640×640,batch=1),mAP@0.5在COCO-val上达38.1%,比同参数量级的YOLO-Nano高4.2个百分点。

更重要的是:它原生支持自定义数据集端到端训练+一键ONNX导出+无依赖推理封装。没有中间转换脚本,没有手动修改导出配置,也没有“导出后结果对不上”的玄学问题。整套流程我们今天就用真实代码走一遍——从准备数据、修改类别、训练模型,到生成ONNX、验证输出,全部可复制、可调试、不跳坑。

2. 环境准备与项目初始化

2.1 硬件与基础环境要求

EagleEye对硬件友好,但为保障训练稳定性与导出一致性,建议按以下配置准备:

  • GPU:NVIDIA RTX 3090 / 4090(显存 ≥24GB,训练阶段需双卡更佳)
  • 系统:Ubuntu 22.04 LTS(推荐,CUDA兼容性最佳)
  • 驱动与工具链
    nvidia-driver-535 cuda-toolkit-12.1 cudnn-8.9.2

注意:不要用conda安装PyTorch——EagleEye官方训练脚本依赖torch.compiletorch.export的特定行为,仅支持pip安装的官方CUDA版本。执行以下命令一次性配齐:

pip3 install torch==2.1.1+cu121 torchvision==0.16.1+cu121 --extra-index-url https://download.pytorch.org/whl/cu121 pip3 install opencv-python==4.8.1.78 numpy==1.24.4 tqdm==4.66.1 onnx==1.15.0 onnxruntime-gpu==1.17.1

2.2 获取EagleEye源码与预训练权重

EagleEye开源在GitHub(非镜像仓库),使用git lfs管理大文件,务必提前安装:

git lfs install git clone https://github.com/alibaba/EagleEye.git cd EagleEye

预训练权重已托管在阿里云OSS,直接下载解压即可:

wget https://eagleeye-models.oss-cn-hangzhou.aliyuncs.com/weights/tinynas_damoyolo_s.pth mkdir -p weights/ mv tinynas_damoyolo_s.pth weights/

此时目录结构应为:

EagleEye/ ├── configs/ │ └── tinynas_damoyolo_s.py # 主配置文件(含类别数、输入尺寸、NAS结构定义) ├── datasets/ │ └── custom/ # 我们将在此创建自定义数据集 ├── models/ ├── tools/ │ ├── train.py # 训练入口 │ ├── export_onnx.py # ONNX导出脚本(本文核心!) │ └── infer_onnx.py # ONNX推理验证脚本 ├── weights/ │ └── tinynas_damoyolo_s.pth

3. 自定义数据集准备与类别配置

3.1 数据集格式:纯正YOLOv5-style,零学习成本

EagleEye完全兼容YOLOv5标准格式,无需额外标注工具或格式转换。你只需组织好以下结构:

datasets/custom/ ├── images/ │ ├── train/ │ │ ├── img_001.jpg │ │ └── img_002.jpg │ └── val/ │ ├── img_011.jpg │ └── img_012.jpg └── labels/ ├── train/ │ ├── img_001.txt # 每行: class_id center_x center_y width height (归一化) │ └── img_002.txt └── val/ ├── img_011.txt └── img_012.txt

关键细节

  • center_x,center_y,width,height均为相对于图像宽高的归一化值(0~1)
  • class_id从0开始编号,对应你后续定义的类别顺序
  • 图像尺寸不限(训练时自动resize),但建议统一为长边≤1920,避免显存溢出

3.2 修改配置文件:只改3处,不碰NAS结构

打开configs/tinynas_damoyolo_s.py,定位以下三处并修改(其他参数保持默认):

# === 1. 修改数据集路径 === data_root = 'datasets/custom/' # ← 指向你的数据集根目录 # === 2. 修改类别数量与名称 === num_classes = 3 # ← 改为你实际类别的数量(例:person, car, traffic_light) classes = ('person', 'car', 'traffic_light') # ← 严格按顺序填写,不可空格/特殊字符 # === 3. 修改训练数据加载器中的类别映射 === train_dataloader = dict( dataset=dict( data_root=data_root, ann_file='labels/train.txt', # ← 若你用txt列表而非文件夹遍历,可指定 data_prefix=dict(img='images/train/'), metainfo=dict(classes=classes), # ← 关键!必须显式传入 ) )

小技巧:如果你的数据集是COCO或Pascal VOC格式,用tools/dataset_converters/下的转换脚本一键转YOLO格式,5分钟搞定。

4. 启动训练:单卡/双卡灵活切换

EagleEye支持torchrun多卡训练,也支持单卡直跑。我们以双RTX 4090训练为例(显存充足,收敛更快):

# 启动双卡训练(自动启用DDP) torchrun --nproc_per_node=2 tools/train.py \ --config configs/tinynas_damoyolo_s.py \ --work-dir work_dirs/custom_tinynas_s \ --amp # 启用混合精度,提速30%且不掉点

训练过程会实时打印:

  • 当前epoch与iter
  • loss_cls(分类损失)、loss_box(框回归损失)、loss_dfl(分布焦点损失)
  • mAP@0.5(验证集指标,每5个epoch计算一次)

避坑提醒

  • 若显存不足报OOM,请在配置中调小train_dataloader.batch_size(默认16,可降至8或4)
  • 不要修改model.backbone.arch字段——这是TinyNAS搜索出的最优结构,人工改动会破坏毫秒级延迟保障

训练约12小时(100 epochs)后,最佳权重将保存在:
work_dirs/custom_tinynas_s/weights/best_ckpt.pth

5. ONNX导出:一行命令,三重验证

EagleEye的ONNX导出不是“能导出就行”,而是保证输入输出语义一致、后处理逻辑内嵌、坐标格式标准统一。执行以下命令:

python tools/export_onnx.py \ --config configs/tinynas_damoyolo_s.py \ --checkpoint work_dirs/custom_tinynas_s/weights/best_ckpt.pth \ --input-shape 1 3 640 640 \ --output-file eagleeye_custom.onnx \ --dynamic-batch \ --simplify # 启用onnxsim优化,减小体积35%

该命令完成三件事:

  1. 加载模型并设置为eval模式,冻结BN层
  2. 构造虚拟输入(torch.randn(1,3,640,640)),触发torch.export导出流程
  3. 内置后处理(NMS + 置信度过滤)直接编译进ONNX图,输出即为最终检测结果

导出成功后,你会得到:

  • eagleeye_custom.onnx(约28MB,含NMS)
  • eagleeye_custom.onnx.simplified(约18MB,已优化)

5.1 验证ONNX输出是否正确?

别急着部署,先用infer_onnx.py做三重校验:

python tools/infer_onnx.py \ --model eagleeye_custom.onnx \ --img-path datasets/custom/images/val/img_011.jpg \ --conf-thres 0.4 \ --iou-thres 0.5 \ --save-img # 生成带框的result.jpg供肉眼比对

校验通过标准

  • result.jpg中的检测框位置、类别、置信度,与原始PyTorch模型tools/infer.py输出完全一致(像素级对齐)
  • ONNX Runtime推理耗时 ≤ PyTorch推理耗时 × 1.05(实测17.1ms vs 17.3ms)
  • 输出tensor shape为[1, N, 6],其中每行=[x1,y1,x2,y2,conf,class_id](标准YOLO格式)

深入看一眼ONNX结构
用Netron打开eagleeye_custom.onnx,你会看到:

  • 输入名:images(shape: [B,3,640,640])
  • 输出名:detections(shape: [B,N,6])
  • 中间无任何Resize/Pad节点——所有预处理(归一化、letterbox)需在ONNX外部完成,这是EagleEye设计的明确约定。

6. 实战部署:Python + C++ 双路径接入

6.1 Python端:3行代码完成推理

import cv2 import numpy as np import onnxruntime as ort # 1. 加载ONNX模型(GPU加速) session = ort.InferenceSession("eagleeye_custom.onnx", providers=['CUDAExecutionProvider']) # 2. 图像预处理(必须!letterbox + 归一化) img = cv2.imread("test.jpg") img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img_resized = cv2.resize(img_rgb, (640, 640)) img_norm = (img_resized.astype(np.float32) / 255.0).transpose(2, 0, 1)[None] # [1,3,640,640] # 3. 推理 & 解析 outputs = session.run(None, {"images": img_norm})[0] # [1,N,6] for det in outputs[0]: # 遍历每个检测结果 x1, y1, x2, y2, conf, cls_id = det if conf > 0.4: cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (0,255,0), 2) cv2.putText(img, f"{classes[int(cls_id)]} {conf:.2f}", (int(x1), int(y1)-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2) cv2.imwrite("result.jpg", img)

6.2 C++端:跨平台集成(Linux/Windows)

EagleEye提供C++推理示例(demo/cpp_inference/),基于ONNX Runtime C API,编译后二进制仅12MB,无Python依赖:

// 示例:加载模型 + 单图推理(省略错误检查) Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "EagleEye"}; Ort::SessionOptions session_options; session_options.SetIntraOpNumThreads(4); session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED); Ort::Session session{env, L"eagleeye_custom.onnx", session_options}; std::vector<const char*> input_names = {"images"}; std::vector<const char*> output_names = {"detections"}; // ...(图像读取、预处理、内存拷贝) auto output_tensors = session.Run(Ort::RunOptions{nullptr}, input_names.data(), &input_tensor, 1, output_names.data(), 1); // output_tensors[0].GetTensorMutableData<float>() 即为 [N,6] 结果

编译命令(Ubuntu):

g++ -std=c++17 demo/cpp_inference/infer.cpp -I/usr/include/onnxruntime \ -L/usr/lib -lonnxruntime -lopencv_core -lopencv_imgproc -lopencv_highgui \ -o eagleeye_cpp

7. 性能对比与落地建议

我们用同一台双4090机器,对比EagleEye与其他主流轻量模型在自定义数据集(3类,2000张图)上的表现:

模型参数量(M)推理延迟(ms)mAP@0.5ONNX导出稳定性是否需后处理
EagleEye (TinyNAS-S)2.117.372.4一次导出,开箱即用❌ 内置NMS
YOLOv8n3.221.869.1导出后需手动加NMS外部实现
NanoDet-m0.914.263.8❌ 导出后坐标错乱频发外部实现
PP-YOLOE-s4.825.671.2稳定❌ 内置

给工程团队的落地建议

  • 选型阶段:若延迟敏感(<20ms)且类别≤10,优先选EagleEye TinyNAS-S;若需更高精度(>75mAP)且算力充足,可试TinyNAS-M(参数量4.3M,延迟24ms)
  • 数据层面:EagleEye对小目标(<32×32)检测鲁棒性强,但建议训练时开启mosaic=0.5增强,避免过拟合
  • 部署边界:ONNX模型不包含图像解码/预处理,务必在业务层统一实现letterbox(保持宽高比)与归一化,否则框坐标偏移

获取更多AI镜像

想探索更多AI镜像和应用场景?访问 CSDN星图镜像广场,提供丰富的预置镜像,覆盖大模型推理、图像生成、视频生成、模型微调等多个领域,支持一键部署。

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.mzph.cn/news/1221876.shtml

如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

企业广告配音新方案:IndexTTS 2.0批量生成实践

企业广告配音新方案&#xff1a;IndexTTS 2.0批量生成实践 你有没有经历过这样的加班夜&#xff1a;市场部刚发来12条新品广告文案&#xff0c;要求明天一早全部配好音——男声沉稳版、女声亲和版、年轻活力版、方言本地化版……还要严格卡在15秒内&#xff0c;不能快半秒&…

2026年合肥专业自助入住系统公司深度评测与推荐

随着数字化转型浪潮席卷酒店住宿业,自助入住系统已从提升效率的辅助工具,演变为重塑客户体验、驱动运营增长的核心技术引擎。尤其在合肥这座快速发展的新一线城市,酒店、公寓、民宿业主正面临着人力成本攀升、客户对…

破解教材下载难题:tchMaterial-parser全功能解析

破解教材下载难题&#xff1a;tchMaterial-parser全功能解析 【免费下载链接】tchMaterial-parser 国家中小学智慧教育平台 电子课本下载工具 项目地址: https://gitcode.com/GitHub_Trending/tc/tchMaterial-parser 教育资源下载工具、电子教材管理、PDF解析器——tchM…

如何在macOS上运行Windows程序:Whisky的跨平台解决方案

如何在macOS上运行Windows程序&#xff1a;Whisky的跨平台解决方案 【免费下载链接】Whisky A modern Wine wrapper for macOS built with SwiftUI 项目地址: https://gitcode.com/gh_mirrors/wh/Whisky 副标题&#xff1a;基于SwiftUI的现代Wine封装工具&#xff0c;让…

LyricsX与iTunes歌词同步解决方案:自动完善音乐库的实用指南

LyricsX与iTunes歌词同步解决方案&#xff1a;自动完善音乐库的实用指南 【免费下载链接】LyricsX &#x1f3b6; Ultimate lyrics app for macOS. 项目地址: https://gitcode.com/gh_mirrors/ly/LyricsX 在数字音乐时代&#xff0c;完整的歌词体验已成为音乐欣赏不可或…

6大防护策略!用luci-app-access-control构建家庭网络安全防线

6大防护策略&#xff01;用luci-app-access-control构建家庭网络安全防线 【免费下载链接】luci-access-control OpenWrt internet access scheduler 项目地址: https://gitcode.com/gh_mirrors/lu/luci-access-control 随着智能家居设备普及&#xff0c;家庭网络边界逐…

VibeThinker-1.5B效率翻倍:优化推理速度的小技巧

VibeThinker-1.5B效率翻倍&#xff1a;优化推理速度的小技巧 在大模型部署动辄需要多卡A100、显存占用动辄20GB以上的今天&#xff0c;一个仅需单张T4&#xff08;甚至RTX 3060&#xff09;就能跑通、显存峰值稳定在1.8GB以内、却能在AIME数学竞赛题和LeetCode Hard算法题上稳…

分子动力学分析实战指南:从认知基础到效能提升的完整路径

分子动力学分析实战指南&#xff1a;从认知基础到效能提升的完整路径 【免费下载链接】mdanalysis MDAnalysis is a Python library to analyze molecular dynamics simulations. 项目地址: https://gitcode.com/gh_mirrors/md/mdanalysis 在分子动力学研究领域&#xf…

告别手动点击!用Open-AutoGLM快速搭建手机AI助手

告别手动点击&#xff01;用Open-AutoGLM快速搭建手机AI助手 你有没有过这样的时刻&#xff1a; 想查个快递&#xff0c;却要解锁、找App、点开、输入单号、等加载…… 想关注一个博主&#xff0c;得打开抖音、搜索ID、点进主页、再点关注——手指划了五下&#xff0c;才完成一…

keil编译器下载v5.06下CAN总线控制程序设计实战案例

以下是对您提供的博文内容进行 深度润色与结构重构后的技术文章 。我以一位深耕嵌入式系统多年、长期使用 Keil STM32 开发工业级 CAN 应用的工程师视角&#xff0c;重写了全文——去除了所有模板化标题、AI腔调和空泛表述&#xff0c;强化了 真实开发语境中的逻辑流、踩坑…

all-MiniLM-L6-v2参数详解:hidden_size=384, num_layers=6, dropout=0.1全解读

all-MiniLM-L6-v2参数详解&#xff1a;hidden_size384, num_layers6, dropout0.1全解读 1. 模型本质&#xff1a;不是“小号BERT”&#xff0c;而是专为嵌入而生的精炼架构 很多人第一眼看到 all-MiniLM-L6-v2&#xff0c;会下意识把它当成“缩水版BERT”——层数少了、维度低…

2步解锁唇语黑科技:让无声交流秒变文字

2步解锁唇语黑科技&#xff1a;让无声交流秒变文字 【免费下载链接】chaplin A real-time silent speech recognition tool. 项目地址: https://gitcode.com/gh_mirrors/chapl/chaplin 当静音成为刚需&#xff0c;如何让唇语成为新的交互语言&#xff1f; 在图书馆想记…

零编码经验?照样用BSHM镜像玩转AI图像处理

零编码经验&#xff1f;照样用BSHM镜像玩转AI图像处理 你是不是也遇到过这些场景&#xff1a; 想给朋友圈照片换个星空背景&#xff0c;却卡在Photoshop的图层蒙版里&#xff1b; 电商运营要批量处理上百张模特图&#xff0c;手动抠图一天都干不完&#xff1b; 设计师赶着交稿…

解锁创意卡牌设计:专业级三国杀武将创作全指南

解锁创意卡牌设计&#xff1a;专业级三国杀武将创作全指南 【免费下载链接】Lyciumaker 在线三国杀卡牌制作器 项目地址: https://gitcode.com/gh_mirrors/ly/Lyciumaker 作为卡牌设计师&#xff0c;我们深知一张成功的三国杀卡牌不仅是视觉艺术的呈现&#xff0c;更是游…

7个技巧打造专业级三国杀卡牌:从新手到高手的创作指南

7个技巧打造专业级三国杀卡牌&#xff1a;从新手到高手的创作指南 【免费下载链接】Lyciumaker 在线三国杀卡牌制作器 项目地址: https://gitcode.com/gh_mirrors/ly/Lyciumaker 一、卡牌制作常见问题解析 1.1 设计痛点与解决方案 传统卡牌制作面临三大核心难题&#…

如何实现VRChat无障碍沟通?VRCT全场景应用指南

如何实现VRChat无障碍沟通&#xff1f;VRCT全场景应用指南 【免费下载链接】VRCT VRCT(VRChat Chatbox Translator & Transcription) 项目地址: https://gitcode.com/gh_mirrors/vr/VRCT 如何突破语言壁垒&#xff1f;VRCT的核心价值定位 在全球化的虚拟社交平台VR…

7个系统焕新技巧:用Dism++解决Windows卡顿的终极优化方案

7个系统焕新技巧&#xff1a;用Dism解决Windows卡顿的终极优化方案 【免费下载链接】Dism-Multi-language Dism Multi-language Support & BUG Report 项目地址: https://gitcode.com/gh_mirrors/di/Dism-Multi-language 系统运行缓慢、磁盘空间告急、更新频繁失败—…

如何用ms-swift实现7B模型4-bit量化?实测分享

如何用ms-swift实现7B模型4-bit量化&#xff1f;实测分享 你是否也遇到过这样的困境&#xff1a;手头只有一张RTX 3090&#xff08;24GB显存&#xff09;&#xff0c;却想跑通Qwen2.5-7B这类主流大模型的微调与部署&#xff1f;下载完模型权重就卡在显存不足&#xff0c;量化脚…

系统化岛屿设计:从空白画布到梦幻乐园的进阶指南

系统化岛屿设计&#xff1a;从空白画布到梦幻乐园的进阶指南 【免费下载链接】HappyIslandDesigner "Happy Island Designer (Alpha)"&#xff0c;是一个在线工具&#xff0c;它允许用户设计和定制自己的岛屿。这个工具是受游戏《动物森友会》(Animal Crossing)启发而…

Scanner类方法项目应用快速上手

以下是对您提供的博文内容进行 深度润色与重构后的技术文章 。我以一位深耕嵌入式Java开发十余年、常年在工业现场调试设备的工程师视角&#xff0c;重新组织逻辑、删减冗余术语、强化工程语感&#xff0c;并彻底去除AI写作痕迹——全文无“本文将…”“综上所述”等模板化表…