Модель детектирования RT-DETR и ее обучение

RT-DETR расшифровывается как Real-Time DEtection TRansformer. В свой работе авторы заявляют, что она эффективнее, чем модели YOLO и при этом лицензия Apache 2.0 позволяет ее использовать в коммерческих целях.

Ссылка на работу: https://arxiv.org/abs/2304.08069

Ссылка на репозиторий: https://github.com/lyuwenyu/RT-DETR

Графики, приведенные авторами:

2024-07-23_16-51-29

Графики вызывают некоторое сомнение по быстродействию. Например, YOLOv8s имеет 11.2 миллионов параметров, а RT-DETR-R18 - 20 миллионов параметров. Но при этом RT-DETR-R18 быстрее. Конечно там другая архитектура, но количество параметров есть количество параметров. Во всяком случае, если оставить вопрос о batch_size - было ли оно справедливым для тестов на всех моделях - то даже в этом случае, если запускать модели на CPU, то очевидно, что 11 миллионов параметров будут вычислены быстрее 20.

Но вернемся к вопросам тренировки модели

 Установка пакетов и подготовка выборки

Для тренировки была использована инфраструктура Kaggle и своя выборка для 6 классов объектов, используемая в одном из проектов, в формате YOLO.

Установка необходимого для модели:

!git clone https://github.com/lyuwenyu/RT-DETR -q

 

%cd RT-DETR

 

%cd rtdetr_pytorch

 

!pip install -r requirements.txt

Перевод выборки из YOLO формата в COCO:

Создание каталогов:

!mkdir /kaggle/working/dataset/

!mkdir /kaggle/working/dataset/images

!mkdir /kaggle/working/dataset/images/train2017

!mkdir /kaggle/working/dataset/images/val2017

!mkdir /kaggle/working/dataset/images/annotations

Установка пакета для конвертации :

!pip install globox

 

Копирование изображений:

!cp /kaggle/input/dataset6/train/*.jpg /kaggle/working/dataset/images/train2017

 

!cp /kaggle/input/dataset6/validation/*.jpg /kaggle/working/dataset/images/val2017

Формирование json файлов COCO:

from globox import *

from pathlib import Path

  

def convert(path,names_file,save_file):

  annotations = AnnotationSet.from_yolo_v5(

    folder=path,

    image_folder=path)

  annotations.save_coco(save_file,auto_ids=True)

 

convert('/kaggle/input/dataset6/train/','classes.names','/kaggle/working/dataset/images/annotations/train.json')

 

convert('/kaggle/input/dataset6/validation/','classes.names','/kaggle/working/dataset/images/annotations/valid.json')

Настройка параметров обучения и тренировка

Создаем свой файл оптимизатора и задаем количество эпох:

optimizer ='''use_ema: True

ema:

  type: ModelEMA

  decay: 0.9999

  warmups: 2000

 

 

find_unused_parameters: True

 

epoches: 50

clip_max_norm: 0.1

 

optimizer:

  type: AdamW

  params:

    -

      params: 'backbone'

      lr: 0.00001

    -

      params: '^(?=.*encoder(?=.*bias|.*norm.*weight)).*$'

      weight_decay: 0.

    -

      params: '^(?=.*decoder(?=.*bias|.*norm.*weight)).*$'

      weight_decay: 0.

 

  lr: 0.0001

  betas: [0.9, 0.999]

  weight_decay: 0.0001

 

 

lr_scheduler:

  type: MultiStepLR

  milestones: [1000]

  gamma: 0.1

'''

pathoptimizer = "/kaggle/working/RT-DETR/rtdetr_pytorch/configs/rtdetr/include/myoptimizer.yml"

 

with open(pathoptimizer, 'w') as f:

    f.write(optimizer)

Создаем свой файл датасета, задаем количество объектов и пути к выборке:

datasetyml='''

task: detection

 

num_classes: 6

remap_mscoco_category: False

 

train_dataloader:

  type: DataLoader

  dataset:

    type: CocoDetection

    img_folder: /kaggle/working/dataset/images/train2017

    ann_file: /kaggle/working/dataset/images/annotations/train.json

    transforms:

      type: Compose

      ops: ~

  shuffle: True

  batch_size: 8

  num_workers: 4

  drop_last: True

 

 

val_dataloader:

  type: DataLoader

  dataset:

    type: CocoDetection

    img_folder: /kaggle/working/dataset/images/val2017

    ann_file: /kaggle/working/dataset/images/annotations/valid.json

    transforms:

      type: Compose

      ops: ~

 

  shuffle: False

  batch_size: 8

  num_workers: 4

  drop_last: False

'''

pathdataset = "/kaggle/working/RT-DETR/rtdetr_pytorch/configs/dataset/mydataset.yml"

with open(pathdataset, 'w') as f:

    f.write(datasetyml)

Создаем свой файл загрузчика, где можем задавать аугментацию:

dataloader="""train_dataloader:

  dataset:

    return_masks: False

    transforms:

      ops:

        - {type: RandomPhotometricDistort, p: 0.5}

        - {type: RandomZoomOut, fill: 0}

        - {type: RandomIoUCrop, p: 0.5}

        - {type: SanitizeBoundingBox, min_size: 1}

        - {type: RandomHorizontalFlip, p: 0.5}

        - {type: Resize, size: [640, 640], }

        # - {type: Resize, size: 639, max_size: 640}

        # - {type: PadToSize, spatial_size: 640}

        - {type: ToImageTensor}

        - {type: ConvertDtype}

        - {type: SanitizeBoundingBox, min_size: 1}

        - {type: ConvertBox, out_fmt: 'cxcywh', normalize: True}

  shuffle: True

  batch_size: 4

  num_workers: 4

  collate_fn: default_collate_fn

 

 

val_dataloader:

  dataset:

    transforms:

      ops:

        # - {type: Resize, size: 639, max_size: 640}

        # - {type: PadToSize, spatial_size: 640}

        - {type: Resize, size: [640, 640]}

        - {type: ToImageTensor}

        - {type: ConvertDtype}

  shuffle: False

  batch_size: 8

  num_workers: 4

  collate_fn: default_collate_fn

"""

pathdataloader = "/kaggle/working/RT-DETR/rtdetr_pytorch/configs/rtdetr/include/mydataloader.yml"

 

with open(pathdataloader, 'w') as f:

    f.write(dataloader)

Создаем основной конфигурационный файл:

config = '''

 

__include__: [

  '../dataset/mydataset.yml',

  '../runtime.yml',

  './include/mydataloader.yml',

  './include/myoptimizer.yml',

  './include/rtdetr_r50vd.yml',

]

  

output_dir: ./output/rtdetr_r18vd_6x_coco

 

PResNet:

  depth: 18

  freeze_at: -1

  freeze_norm: False

  pretrained: True

 

HybridEncoder:

  in_channels: [128, 256, 512]

  hidden_dim: 256

  expansion: 0.5

  

RTDETRTransformer:

  eval_idx: -1

  num_decoder_layers: 3

  num_denoising: 100

  

optimizer:

  type: AdamW

  params:

    -

      params: '^(?=.*backbone)(?=.*norm).*$'

      lr: 0.00001

      weight_decay: 0.

    -

      params: '^(?=.*backbone)(?!.*norm).*$'

      lr: 0.00001

    -

      params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bias)).*$'

      weight_decay: 0.

 

  lr: 0.0001

  betas: [0.9, 0.999]

  weight_decay: 0.0001

'''

pathconfig = "/kaggle/working/RT-DETR/rtdetr_pytorch/configs/rtdetr/myrtdetr_r18vd_6x_coco.yml"

with open(pathconfig, 'w') as f:

    f.write(config)

Запускаем обучение:

!export CUDA_VISIBLE_DEVICES=0

 

!python tools/train.py -c configs/rtdetr/myrtdetr_r18vd_6x_coco.yml

Конечный результат:

IoU metric: bbox

 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.569

 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.754

 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.625

 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.401

 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.434

 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.664

 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.435

 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.731

 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.805

 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.669

 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.761

 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.863

best_stat:  {'epoch': 49, 'coco_eval_bbox': 0.5687755125856716}

Для этой выборки результат прилично хуже, чем для YOLOX small и YOLOv5 small. Но я допускаю, что тут нужно еще поработать с аугментацией для этой модели и возможно, что на других данных результат будет действительно лучше, но пока есть то, что есть.

 ONNX и TensorRT

Перевод в ONNX:

!python tools/export_onnx.py -c configs/rtdetr/myrtdetr_r18vd_6x_coco.yml -r output/rtdetr_r18vd_6x_coco/checkpoint0049.pth --check

 

В TensorRT можно перевести на другом, нужном для нас компьютере:

trtexec --onnx=modelRTDETR.onnx --workspace=4096 --shapes=images:1x3x640x640 --saveEngine=rtdetr_r50vd_6x_coco.trt --avgRuns=100 --fp16

Модель переводится версией 8.5.3.1 без проблем, хоть у нее и динамические параметры.

2024-07-23_17-23-39

И в этом ее положительное отличие от другой моделей

pp-picodet и rtmdet

где нужно часть слоев удалять, восстанавливая их код на CPU.