RT-DETR¶

This page describes how to use the official RT-DETR implementation with LightlyTrain.

Note

RT-DETR is not a pip-installable Python package. For this reason, RT-DETR is not fully integrated with LightlyTrain and has to be installed manually.

Supported Models¶

The following RT-DETR model variants are supported:

rtdetr_r18vd
rtdetr_r34vd
rtdetr_r50vd
rtdetr_r50vd_m
rtdetr_r101vd

RT-DETR Installation¶

To install RT-DETR for use with LightlyTrain, first clone the repository:

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

We assume that all commands are executed from the RT-DETR/rtdetr_pytorch directory unless otherwise stated:

cd RT-DETR/rtdetr_pytorch

Next, create a Python environment using your preferred tool and install the required dependencies. Some dependencies are fixed to specific versions to ensure compatibility with RT-DETR:

pip install lightly-train -r requirements.txt "numpy==1.26" "pytorch-lightning==2.1.0" pycocotools

Pretrain an RT-DETR Model¶

Run the following script to pretrain an RT-DETR model. The script must either be executed from inside the RT-DETR/rtdetr_pytorch directory or the RT-DETR/rtdetr_pytorch directory must be added to the Python path.

# pretrain_rtdetr.py
from typing import Dict

import torch
from torch import Tensor
from torch.nn import AdaptiveAvgPool2d, Module

import lightly_train
from src.core import YAMLConfig


class RTDETRModelWrapper(Module):
    def __init__(self, model: Module):
        super().__init__()
        self._model = model
        self._pool = AdaptiveAvgPool2d((1, 1))

    def forward_features(self, x: Tensor) -> Dict[str, Tensor]:
        features = self._model.backbone(x)[-1]
        return {"features": features}

    def forward_pool(self, x: Dict[str, Tensor]) -> Dict[str, Tensor]:
        return {"pooled_features": self._pool(x["features"])}

    def feature_dim(self) -> int:
        return self._model.backbone.out_channels[-1]


if __name__ == "__main__":
    # Load the RT-DETR model
    # Change the config file to the desired RT-DETR model
    config = YAMLConfig("configs/rtdetr/rtdetr_r50vd_6x_coco.yml")
    model = config.model

    # Pretrain the model
    wrapped_model = RTDETRModelWrapper(model)
    lightly_train.train(
        out="out/my_experiment",
        model=wrapped_model,
        data="./dataset/coco/train2017", # Replace with your dataset path.
    )

    # Save the pretrained model for fine-tuning
    torch.save(model.state_dict(), "out/my_experiment/pretrained_model.pt")

Fine-Tune an RT-DETR Model¶

After the pretraining completes, the model can be fine-tuned using the default RT-DETR training script by providing the path to the pretrained model:

# Training on single-gpu
export CUDA_VISIBLE_DEVICES=0
python tools/train.py -c configs/rtdetr/rtdetr_r50vd_6x_coco.yml --resume out/my_experiment/pretrained_model.pt

See the RT-DETR repository for more information on how to fine-tune a model.