LightlyTrain with RF-DETR

In this notebook we will demonstrate how you can use LightlyTrain to pretrain an RF-DETR model from Roboflow. For now, rfdetr only supports training with datasets in COCO JSON format. To this end, for pretraining we use the raw images (no labels) from the COCO-minitrain dataset, a subset of the COCO dataset with 25k images, and for fine-tuning we use the Roboflow’s Coconut Custom Dataset.

Open In Colab

Important: When running on Google Colab make sure to select a GPU runtime for faster processing. You can do this by going to Runtime > Change runtime type and selecting a GPU hardware accelerator.

Installation

You can install lightly_train directly from PyPI using pip with support of rfdetr package.

!pip install "lightly-train[rfdetr]"

Important: LightlyTrain is officially supported on

  • Linux: CPU or CUDA

  • MacOS: CPU only

  • Windows (experimental): CPU or CUDA

We are planning to support MPS for MacOS.

Check the installation instructions for more details on installation.

Pretrain on COCO-minitrain-25k Dataset

We use the COCO-minitrain dataset, a subset of the COCO dataset with 25k images for pretraining the RF-DETR model.

Download the Dataset

We can download the COCO-minitrain dataset (25k images) directly from HuggingFace

!wget https://huggingface.co/datasets/bryanbocao/coco_minitrain/resolve/main/coco_minitrain_25k.zip

… unzip it…

!unzip coco_minitrain_25k.zip

… and since LightlyTrain does not require any labels, we can also confidently delete all the labels:

!rm -rf coco_minitrain_25k/labels

Pretrain an RF-DETR Model

Pretraining an RF-DETR model with LightlyTrain is straightforward:

import lightly_train

if __name__ == "__main__":
    lightly_train.pretrain(
        out="out/my_experiment",  # Output directory.
        data="coco_minitrain_25k/images",  # Directory with images.
        model="rfdetr/rf-detr-base",  # Pass the RF-DETR model.
        epochs=5,  # Number of epochs to train
        batch_size=16,  # Batch size
        overwrite=True,
    )

Fine-tune on Coconuts Custom Dataset

We use Roboflow’s Coconut Custom Dataset for fine-tuning.

Download the Dataset

The dataset can be directly downloaded via Roboflow API:

from roboflow import Roboflow

rf = Roboflow(api_key="your_roboflow_api_key")

project = rf.workspace("ravi-mgvlz").project("coconut-custom-dataset")
version = project.version(3)
finetune_dataset = version.download("coco")

Fine-tune an RF-DETR Model

You can directly use the rfdetr package for fine-tuning.

from rfdetr import RFDETRBase

model = RFDETRBase(
    pretrain_weights="out/my_experiment/exported_models/exported_last.pt"
)
model.train(dataset_dir=finetune_dataset.location, epochs=10, batch_size=4, lr=1e-4)