Skip to content

Latest commit

 

History

History
139 lines (106 loc) · 5.5 KB

README_en.md

File metadata and controls

139 lines (106 loc) · 5.5 KB

RT-DETR

  • Train PaddleDetection's RT-DETR model with a custom dataset for onnx and TensorRT deployment without licensing issues.

Training

1. Select a Pretrained Model for Fine-tuning

Model Epoch Backbone Input Shape $AP^{val}$ $AP^{val}_{50}$ Params (M) FLOPs (G) T4 TensorRT FP16 (FPS) Pretrained Model Config
RT-DETR-R18 6x ResNet-18 640 46.5 63.8 20 60 217 download config
RT-DETR-R34 6x ResNet-34 640 48.9 66.8 31 92 161 download config
RT-DETR-R50-m 6x ResNet-50 640 51.3 69.6 36 100 145 download config
RT-DETR-R50 6x ResNet-50 640 53.1 71.3 42 136 108 download config
RT-DETR-R101 6x ResNet-101 640 54.3 72.7 76 259 74 download config
RT-DETR-L 6x HGNetv2 640 53.0 71.6 32 110 114 download config
RT-DETR-X 6x HGNetv2 640 54.8 73.1 67 234 74 download config
RT-DETR-H 6x HGNetv2 640 56.3 74.8 123 490 40 download config

2. Update Model Configuration

Locate the corresponding YAML configuration file, e.g., "configs/rtdetr/rtdetr_r50vd_6x_coco.yml".

  • For COCO datasets, set: 
    _BASE_: [ '../datasets/coco_detection.yml' ]
  • For custom datasets, set: 
    _BASE_: [ '../datasets/custom_dataset.yml' ]

Custom Model Training

1. Dataset Preparation

Organize your dataset in the COCO format:

dataset/
├── annotations/
│   ├── instances_train.json
│   ├── instances_val.json
│   ├── instances_test.json
├── train/
│   ├── image1.jpg
│   ├── image2.jpg
├── val/
│   ├── image3.jpg
│   ├── image4.jpg
├── test/
    ├── image5.jpg
    ├── image6.jpg

2. Dataset YAML Configuration

Specify the dataset's location in the YAML file, e.g., src/PaddleDetection/configs/datasets/custom_dataset.yml:

metric: COCO
num_classes: 4

TrainDataset:
    name: COCODataSet
    image_dir: train
    anno_path: annotations/train_datasets.json
    dataset_dir: /rt-detr-paddle-build-onnx-tensorrt/src/PaddleDetection/dataset
    data_fields: ['image', 'gt_bbox', 'gt_class', 'is_crowd']

EvalDataset:
    name: COCODataSet
    image_dir: val
    anno_path: annotations/valid_datasets.json
    dataset_dir: /rt-detr-paddle-build-onnx-tensorrt/src/PaddleDetection/dataset
    allow_empty: true

TestDataset:
    name: COCODataSet
    image_dir: test
    anno_path: annotations/test_datasets.json
    dataset_dir: /rt-detr-paddle-build-onnx-tensorrt/src/PaddleDetection/dataset

3. Start Training

If you wish to log training using wandb, log in with your API key and include --use_wandb True:

wandb login

Train the model:

cd /rt-detr-paddle-build-onnx-tensorrt/src/PaddleDetection
export CUDA_VISIBLE_DEVICES=0,1
python -m paddle.distributed.launch --gpus 0,1 tools/train.py -c configs/rtdetr/rtdetr_r50vd_6x_coco.yml --fleet --eval

The training process will output .pdparams files for the model.

Building ONNX and TensorRT Models

1. ONNX Model Conversion

First, generate .pdmodel and .pdparams files. Then, convert to ONNX format:

bash create_pdmodel.sh
bash create_onnx.sh

2. TensorRT Model Conversion

After creating the ONNX model, build the TensorRT model:

bash create_tensorrt.sh

3. Testing the TensorRT Model

Run a prediction test on the TensorRT model:

python utils/trt_predict.py
Input Data Result Data
Input Result