Deformable DETR

Deformable DETR is an object-detection model that is included in the TAO Toolkit. It supports the following tasks:

• convert

• train

• evaluate

• inference

• export

These tasks can be invoked from the TAO Toolkit Launcher using the following convention on the command-line:

            

            
tao model deformable_detr <sub_task> <args_per_subtask>
        

where, args_per_subtask are the command-line arguments required for a given subtask. Each subtask is explained in detail in the following sections.

Deformable DETR expects directories of images for training or validation and annotated JSON files in COCO format.

Sharding the Data (Optional)

For a large dataset, you can optionally use convert to shard the dataset into smaller chunks to reduce the memory burden. In this process, KITTI-based annotations are converted into smaller sharded JSON files, similar to other object detection networks. Here is an example spec file for converting KITTI-based folders into multiple sharded JSON files.

            

            
input_source: /workspace/tao-experiments/data/sequence.txt
results_dir: /workspace/tao-experiments/sharded
image_dir_name: images
label_dir_name: labels
num_shards: 32
num_partitions: 1
mapping_path: /path/to/your_category_mapping
        

The details of each parameter are summarized in the table below:

The category mapping should contain mapping of your dataset and be in reverse alphabetical order. The default mapping is shown below:

            

            
DEFAULT_TARGET_CLASS_MAPPING = {
  "Person": "person",
  "Person Group": "person",
  "Rider": "person",
  "backpack": "bag",
  "face": "face",
  "large_bag": "bag",
  "person": "person",
  "person group": "person",
  "person_group": "person",
  "personal_bag": "bag",
  "rider": "person",
  "rolling_bag": "bag",
  "rollingbag": "bag",
  "largebag": "bag",
  "personalbag": "bag"
}
        

The following example shows how to use the command:

            

            
tao model deformable_detr convert -e /path/to/spec.yaml
        

The training experiment spec file for Deformable DETR includes model, train, and dataset parameters. Here is an example spec file for training a Deformable DETR model with a resnet50 backbone on a COCO dataset.

            

            
dataset:
  train_data_sources:
    - image_dir: /path/to/coco/train2017/
      json_file: /path/to/coco/annotations/instances_train2017.json
  val_data_sources:
    - image_dir: /path/to/coco/val2017/
      json_file: /path/to/coco/annotations/instances_val2017.json
  num_classes: 91
  batch_size: 4
  workers: 8
  augmentation:
    scales: [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800]
    input_mean: [0.485, 0.456, 0.406]
    input_std: [0.229, 0.224, 0.225]
    horizontal_flip_prob: 0.5
    train_random_resize: [400, 500, 600]
    train_random_crop_min: 384
    train_random_crop_max: 600
    random_resize_max_size: 1333
    test_random_resize: 800
model:
  pretrained_model_path: /path/to/your-pretrained-backbone-model
  backbone: resnet_50
  train_backbone: True
  num_feature_levels: 4
  dec_layers: 6
  enc_layers: 6
  num_queries: 300
  with_box_refine: True
  dropout_ratio: 0.3
train:
  optim:
    lr_backbone: 2e-5
    lr: 2e-4
    lr_steps: [10, 20, 30, 40]
    momentum: 0.9
  num_epochs: 50
        

model

The model parameter provides options to change the Deformable DETR architecture.

            

            
model:
  pretrained_model_path: /path/to/your-resnet50-pretrained-model
  backbone: resnet_50
  train_backbone: True
  num_feature_levels: 4
  dec_layers: 6
  enc_layers: 6
  num_queries: 300
  with_box_refine: True
  dropout_ratio: 0.3
        

train

The train parameter defines the hyperparameters of the training process.

            

            
train:
  optim:
    lr: 0.0001
    lr_backbone: 0.00001
    momentum: 0.9
    weight_decay: 0.0001
    lr_scheduler: MultiStep
    lr_steps: [10, 20, 30, 40]
    lr_decay: 0.1
  num_epochs: 50
  checkpoint_interval: 1
  precision: fp32
  distributed_strategy: ddp
  activation_checkpoint: True
  num_gpus: 8
  num_nodes: 1
        

optim

The optim parameter defines the config for the optimizer in training, including the learning rate, learning scheduler, and weight decay.

            

            
optim:
  lr: 0.0002
  lr_backbone: 0.00002
  momentum: 0.9
  weight_decay: 0.0001
  lr_scheduler: MultiStep
  lr_steps: [10, 20, 30, 40]
  lr_decay: 0.1
        

dataset

The dataset parameter defines the dataset source, training batch size, and augmentation.

            

            
dataset:
  train_data_sources:
    - image_dir: /path/to/coco/images/train2017/
      json_file: /path/to/coco/annotations/instances_train2017.json
  val_data_sources:
    - image_dir: /path/to/coco/images/val2017/
      json_file: /path/to/coco/annotations/instances_val2017.json
  test_data_sources:
    image_dir: /path/to/coco/images/val2017/
    json_file: /path/to/coco/annotations/instances_val2017.json
  infer_data_sources:
    image_dir: /path/to/coco/images/val2017/
    classmap: /path/to/coco/annotations/coco_classmap.txt
  num_classes: 91
  batch_size: 4
  workers: 8
        

augmentation

The augmentation parameter contains hyperparameters for augmentation.

            

            
augmentation:
  scales: [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800]
  input_mean: [0.485, 0.456, 0.406]
  input_std: [0.229, 0.224, 0.225]
  horizontal_flip_prob: 0.5
  train_random_resize: [400, 500, 600]
  train_random_crop_min: 384
  train_random_crop_max: 600
  random_resize_max_size: 1333
  test_random_resize: 800
        

To train a Deformable DETR model, use this command:

            

            
tao model deformable_detr train [-h] -e <experiment_spec>
                                [-r <results_dir>]
                                [-k <key>]
        

Required Arguments

• -e, --experiment_spec: The experiment specification file to set up the training experiment

Optional Arguments

• -r, --results_dir: The path to the folder where the experiment outputs should be written. If not specified, the results_dir from the spec file will be used.

• -k, --key: A user-specific encoding key to save or load a .tlt model. If not specified, the model checkpoint will not be encrypted.

• --gpus: The number of GPUs to run training

• --num_nodes: The number of nodes to run training. If this value is larger than 1, distributed multi-node training is enabled.

• -h, --help: Show this help message and exit.

Sample Usage

Here’s an example of the train command:

            

            
tao deformable_detr model train -e /path/to/spec.yaml
        

Optimizing Resource for training Deformable DETR

Training Deformable DETR requires strong GPUs (e.g. V100/A100) with at least 15GB of VRAM and a lot of CPU memory to be trained on a standard dataset like COCO. In this section, we outline some of the strategies you can use to launch training with only limited resources.

Optimize GPU Memory

There are various ways to optimize GPU memory usage. One obvious trick is to reduce dataset.batch_size. However, this can cause your training to take longer than usual. Hence, we recommend setting below configurations in order to optimize GPU consumption.

• Set train.precision to fp16 to enable automatic mixed precision training. This can reduce your GPU memory usage by 50%.

• Set train.activation_checkpoint to True to enable activation checkpointing. By recomputing the activations instead of caching them into memory, the memory usage can be improved.

• Set train.distributed_strategy to ddp_sharded to enabled Sharded DDP training. This will share gradient calculation across different processes to help reduce GPU memory.

• Try using more lightweight backbones like gc_vit_xxtiny or freeze the backbone through setting model.train_backbone to False.

• Try changing the augmentation resolution in dataset.augmentation depending on your dataset.

Optimize CPU Memory

To speed up data loading, it is a common practice to set high number of workers to spawn multiple processes. However, this can cause your CPU memory to become Out of Memory if the size of your annotation file is very large. Hence, we recommend setting below configurations in order to optimize CPU consumption.

• Set dataset.dataset_type to serialized so that the COCO-based annotation data can be shared across different subprocesses.

• Set dataset.augmentation.fixed_padding to True so that images are padded before the batch formulation. Due to random resize and random crop augmentation during training, the resulting image resolution after transform can vary across images. Such variable image resolutions can cause memory leak and the CPU memory to slowly stacks up until it becomes Out of Memory in the middle of training. This is the limitation of PyTorch so we advise setting fixed_padding to True to help stablize the CPU memory usage.

evaluate

The evaluate parameter defines the hyperparameters of the evaluate process.

            

            
evaluate:
  checkpoint: /path/to/model.pth
  conf_threshold: 0.0
  num_gpus: 1
        

To run evaluation with a Deformable DETR model, use this command:

            

            
tao model deformable_detr evaluate [-h] -e <experiment_spec>
                                   [-r <results_dir>]
                                   [-k <key>]
                                    evaluate.checkpoint=<model to be evaluated>
        

Required Arguments

• -e, --experiment_spec: The experiment spec file to set up the evaluation experiment.

Optional Arguments

• -k, --key: A user-specific encoding key to save or load a .tlt model. If not specified, a .pth model must be used.

• -r, --results_dir: The directory where the evaluation result will be stored

• evaluate.checkpoint: The .tlt or .pth model to be evaluated

Sample Usage

Here’s an example of using the evaluate command:

            

            
tao model deformable_detr evaluate -e /path/to/spec.yaml -r /path/to/results/ evaluate.checkpoint=/path/to/model.pth
        

inference

The inference parameter defines the hyperparameters of the inference process.

            

            
inference:
  checkpoint: /path/to/model.pth
  conf_threshold: 0.5
  num_gpus: 1
  color_map:
    person: red
    car: blue
        

The inference tool for Deformable DETR models can be used to visualize bboxes and generate frame-by- frame KITTI format labels on a directory of images.

            

            
tao model deformable_detr inference [-h] -e <experiment spec file>
                                    [-r <results_dir>]
                                    [-k <key>]
                                    inference.checkpoint=<model to be inferenced>
        

Required Arguments

• -e, --experiment_spec: The experiment spec file to set up the inference experiment

Optional Arguments

• -k, --key: A user-specific encoding key to save or load a .tlt model. If not specified, a .pth model must be used.

• -r, --results_dir: The directory where the inference result will be stored

• inference.checkpoint: The .tlt or .pth model for inference

Sample Usage

Here’s an example of using the inference command:

            

            
tao model deformable_detr inference -e /path/to/spec.yaml -r /path/to/results/ inference.checkpoint=/path/to/model.pth
        

export

The export parameter defines the hyperparameters for the export process.

            

            
export:
  checkpoint: /path/to/model.pth
  onnx_file: /path/to/model.onnx
  on_cpu: False
  opset_version: 12
  input_channel: 3
  input_width: 960
  input_height: 544
  batch_size: -1
        

            

            
tao model deformable_detr export [-h] -e <experiment spec file>
                                 [-r <results_dir>]
                                 [-k <key>]
                                 export.checkpoint=<model to export>
                                 export.onnx_file=<onnx path>
        

Required Arguments

• -e, --experiment_spec: The path to an experiment spec file

Optional Arguments

• -k, --key: A user-specific encoding key to save or load a .tlt model. If not specified, a .pth model must be used.

• -r, --results_dir: The directory where the inference result is stored

• export.checkpoint: The .tlt or .pth model to export

• export.onnx_file: The path where the .etlt or .onnx model will be saved

Sample Usage

Here’s an example of using the export command:

            

            
tao model deformable_detr export -e /path/to/spec.yaml export.checkpoint=/path/to/model.pth export.onnx_file=/path/to/model.onnx
        

For deployment, please refer to TAO Deploy documentation.

Refer to the Integrating a Deformable DETR Model page for more information about deploying a Deformable DETR model to DeepStream.