By Toby2718 in 객체인식 — 2025년 3월 31일

RT-DETR 객체인식이란? + 커스텀 트레이닝도 진행해보자!

RT-DETR 객체인식 모델을 소개하고, Ultralytics를 활용한 커스텀 데이터셋 학습 과정을 단계별로 안내합니다. 트랜스포머 기반의 고속 객체 탐지 기술과 yaml 설정, 파이썬 학습 코드까지! labelImg로 라벨링된 YOLO 포맷 데이터셋을 RT-DETR에 적용하는 실습 예제 포함.

객체 인식에 관한 연구를 진행하며, 다양한 객체인식 모델들을 공부하고 사용해 보고있다. 그 중 한가지인 RT-DETR을 오늘 소개해보려고 한다.

먼저 RT-DETR은 바이두에서 개발한 트랜스포머 기반 객체 탐지 모델이다.
ultralytics 라이브러리를 이용하면 간단한 파이썬 코드로 쉽게 학습시킬 수 있으며, yolo 커스텀 트레이닝 데이터셋도 그대로 사용가능하여 무척 편리하다.

모델에 대해서 간단하게 알아보자면

RT-DETR은 DETR 아키텍처를 기반으로 하면서도 실시간 감지를 위해 여러 최적화된 기술을 적용한 것이 특징이다. 대표적으로는 스케일 간 융합(CCFM)과 스케일 내 상호 작용(AIFI) 모듈을 활용하는 ‘효율적인 하이브리드 인코더’를 통해 멀티스케일 특징 정보를 빠르고 정확하게 학습한다. 더불어 IoU 인식 쿼리 선택(IoU-aware Query Selection)을 도입해 중요한 객체 정보를 더욱 정밀하게 포착할 수 있으며, 재학습 없이도 디코더 레이어 수를 유연하게 조절함으로써 추론 속도와 정확도를 상황에 따라 최적화할 수 있다는 장점을 지닌다.

이제 이 모델을 자체 데이터셋로 학습시켜보자.
먼저 데이터셋을 준비해야한다.
데이터셋 준비 과정은 이전글을 참고하기 바란다.

dataset/
│
├── train/           # 학습 데이터 이미지 폴더
│   ├── image1.jpg
│   ├── image2.jpg
│   └── ...
│
├── val/             # 검증 데이터 이미지 폴더
│   ├── image1.jpg
│   ├── image2.jpg
│   └── ...
│
└── labels/          # 라벨 폴더
    ├── train/       # 학습 데이터의 라벨 (텍스트 파일들)
    │   ├── image1.txt
    │   ├── image2.txt
    │   └── ...
    │
    └── val/         # 검증 데이터의 라벨 (텍스트 파일들)
        ├── image1.txt
        ├── image2.txt
        └── ...

여기까지 준비되었다면 yaml파일을 생성해야한다.
yaml파일에는 준비한 데이터셋의 경로와 설정한 클래스명을 작성해주면 된다.
한가지 예시 파일은 아래와 같다.

train: C:\Users\kkwjk2718\Desktop\test\train\imaes # 학습 이미지 경로
val: C:\Users\kkwjk2718\Desktop\test\valid\images # 검증 이미지 경로
test: C:\Users\kkwjk2718\Desktop\test\test\images # 테스트 이미지 경로

names: # 클래스 이름
  0: card
  1: red
  2: box

데이터셋이 모두 준비되었다면, 이제는 정말 학습을 진행할 차례이다.

# run.py
from ultralytics import RTDETR

def main():
    model = RTDETR("rtdetr-l.pt")
    model.train(
        data="rtdetr.yaml",
        epochs=100,
        imgsz=640,
        batch=8,
    )

if __name__ == "__main__":
    main()

위 코드는 에포크 100, 배치사이즈 8정도로 돌린 코드이다. 학습시키는데 gpu리소스를 꽤 소모하므로, 이점 고려하여 학습 진행하기를 바란다.(필자는 4060ti 16gb에서 학습 진행하였다)