ultralytics框架进行RT-DETR目标检测训练

DETR提出以来,其采用匈牙利匹配方式真正的实现了端到端检测效果,避免了NMS等后处理过程,同时,相较CNN的局部特征提取,其凭借着Transformer强大的全局特征提取能力,在目标检测领域可谓大杀四方,基于Transformer的目标检测方法因此层出不穷。

然后,尽管DETR类目标检测方法具有较好的数据拟合能力,但Transformer本身的计算复杂度较高,这使其很难完成实时检测任务,而今天我们则要介绍的便是百度提出的实时DETR目标检测方法,这个方法我已经在先前的博客中有过介绍,当时是基于RT-DETR的源码进行介绍的,今天我们则要介绍的是ultralytics中的RT-DETR模型。

模型结构

RT-DETR模型结构如下:

python 复制代码
                   from  n    params  module                                       arguments                     
  0                  -1  1      9536  ultralytics.nn.modules.block.ResNetLayer     [3, 64, 1, True, 1]           
  1                  -1  1    215808  ultralytics.nn.modules.block.ResNetLayer     [64, 64, 1, False, 3]         
  2                  -1  1   1219584  ultralytics.nn.modules.block.ResNetLayer     [256, 128, 2, False, 4]       
  3                  -1  1   7098368  ultralytics.nn.modules.block.ResNetLayer     [512, 256, 2, False, 6]       
  4                  -1  1  14964736  ultralytics.nn.modules.block.ResNetLayer     [1024, 512, 2, False, 3]      
  5                  -1  1    524800  ultralytics.nn.modules.conv.Conv             [2048, 256, 1, 1, None, 1, 1, False]
  6                  -1  1    789760  ultralytics.nn.modules.transformer.AIFI      [256, 1024, 8]                
  7                  -1  1     66048  ultralytics.nn.modules.conv.Conv             [256, 256, 1, 1]              
  8                  -1  1         0  torch.nn.modules.upsampling.Upsample         [None, 2, 'nearest']          
  9                   3  1    262656  ultralytics.nn.modules.conv.Conv             [1024, 256, 1, 1, None, 1, 1, False]
 10            [-2, -1]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 11                  -1  3   2232320  ultralytics.nn.modules.block.RepC3           [512, 256, 3]                 
 12                  -1  1     66048  ultralytics.nn.modules.conv.Conv             [256, 256, 1, 1]              
 13                  -1  1         0  torch.nn.modules.upsampling.Upsample         [None, 2, 'nearest']          
 14                   2  1    131584  ultralytics.nn.modules.conv.Conv             [512, 256, 1, 1, None, 1, 1, False]
 15            [-2, -1]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 16                  -1  3   2232320  ultralytics.nn.modules.block.RepC3           [512, 256, 3]                 
 17                  -1  1    590336  ultralytics.nn.modules.conv.Conv             [256, 256, 3, 2]              
 18            [-1, 12]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 19                  -1  3   2232320  ultralytics.nn.modules.block.RepC3           [512, 256, 3]                 
 20                  -1  1    590336  ultralytics.nn.modules.conv.Conv             [256, 256, 3, 2]              
 21             [-1, 7]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 22                  -1  3   2232320  ultralytics.nn.modules.block.RepC3           [512, 256, 3]                 
 23        [16, 19, 22]  1   7310072  ultralytics.nn.modules.head.RTDETRDecoder    [4, [256, 256, 256]]          
rtdetr-resnet50 summary: 402 layers, 42,768,952 parameters, 42,768,952 gradients

训练代码

这里我们使用ResNet50作为我们的特征提取网络,从论文的结果来看,RT-DETR只需要24epoch便能达到一个较好的效果。

python 复制代码
from ultralytics import RTDETR
# 加载预训练模型
model = RTDETR("rtdetr-resnet50.yaml")
# 开始训练
results = model.train(
    data="others.yaml",
    epochs=24,
    batch=6,       # 根据GPU显存调整(T4建议batch=8)
    imgsz=640,
    device="0",     # 指定GPU ID
    optimizer="AdamW",
    lr0=1e-4,
    warmup_epochs=4,
    label_smoothing=0.1,
    amp=True
)

相较于YOLO系列的目标检测方法,其训练速度要慢很多,其需要一个epoch需要的时间大概是6分钟,不过从结果来看,其只需要20 epoch便能达到一个较好的效果。

从模型文件大小来看,RT-DETR的文件大小约为245MB,而YOLO模型的大小多在1020MB

最终结果如下:


相关推荐
OAFD.4 小时前
机器学习之线性回归:原理、实现与实践
人工智能·机器学习·线性回归
SHIPKING3936 小时前
【机器学习&深度学习】LMDeploy的分布式推理实现
人工智能·深度学习
mit6.8246 小时前
[RestGPT] docs | RestBench评估 | 配置与环境
人工智能·python
CareyWYR7 小时前
每周AI论文速递(250818-250822)
人工智能
门思科技7 小时前
LoRaWAN 的网络拓扑全解析:架构、原理与应用实践
服务器·网络·人工智能·科技·物联网·架构
兔子的倔强7 小时前
Transformer在文本、图像和点云数据中的应用——经典工作梳理
人工智能·深度学习·transformer
lxmyzzs8 小时前
【图像算法 - 21】慧眼识虫:基于深度学习与OpenCV的农田害虫智能识别系统
人工智能·深度学习·opencv·算法·yolo·目标检测·计算机视觉
Gloria_niki8 小时前
机器学习之K 均值聚类算法
人工智能·机器学习
AI人工智能+8 小时前
表格识别技术:通过图像处理与深度学习,将非结构化表格转化为可编辑结构化数据,推动智能化发展
人工智能·深度学习·ocr·表格识别
深圳多奥智能一卡(码、脸)通系统8 小时前
智能二维码QR\刷IC卡\人脸AI识别梯控系统功能设计需基于模块化架构,整合物联网、生物识别、权限控制等技术,以下是多奥分层次的系统设计框架
人工智能·门禁·电梯门禁·二维码梯控·梯控·电梯