当前位置：首页 > news >正文

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

news 来源：原创 2025/5/7 9:40:10

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

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

模型结构

RT-DETR模型结构如下：

                   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只需要24个epoch便能达到一个较好的效果。

在这里插入图片描述

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 IDoptimizer="AdamW",lr0=1e-4,warmup_epochs=4,label_smoothing=0.1,amp=True
)