CenterPose

CenterPose is a category-level object pose estimation model included in the TAO Toolkit. It supports the following tasks:

• train

• evaluate

• inference

• export

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

            

            
tao model centerpose <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.

CenterPose expects directories of images and annotated JSON files for training or validation. See the CenterPose Data Format page for more information about the input data format.

The training experiment spec file for CenterPose includes model, train, and dataset parameters. Here is an example spec file for training a CenterPose model with a fan_small backbone on a Google Objectron dataset bike category.

            

            
dataset:
  train_data: /path/to/category/train/
  val_data: /path/to/category/val/
  num_classes: 1
  batch_size: 64
  workers: 4
  category: bike
  num_symmetry: 1
  max_objs: 10

train:
  num_gpus: 1
  validation_interval: 20
  checkpoint_interval: 20
  num_epochs: 140
  clip_grad_val: 100.0
  randomseed: 317
  resume_training_checkpoint_path: null
  precision: "fp32"

  optim:
    lr: 6e-05
    lr_steps: [90, 120]

model:
  down_ratio: 4
  use_pretrained: True
  backbone:
    model_type: fan_small
    pretrained_backbone_path: /path/to/your-fan-small-pretrained-model
        

model

The model parameter provides options to change the CenterPose architecture.

            

            
model:
down_ratio: 4
use_pretrained: False
backbone:
  model_type: fan_small
  pretrained_backbone_path: /path/to/your-fan-small-pretrained-model
        

backbone

The backbone parameter provides options to change the CenterPose backbone architecture.

            

            
backbone:
  model_type: fan_small
  pretrained_backbone_path: /path/to/your-fan-small-pretrained-model
        

train

The train parameter defines the hyperparameters of the training process.

            

            
train:
  num_gpus: 1
  validation_interval: 20
  checkpoint_interval: 20
  num_epochs: 140
  clip_grad_val: 100.0
  randomseed: 317
  resume_training_checkpoint_path: null
  precision: "fp32"
  optim:
    lr: 6e-05
    lr_steps: [90, 120]
        

optim

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

            

            
optim:
  lr: 6e-05
  lr_steps: [90, 120]
        

dataset

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

            

            
dataset:
  train_data: /path/to/category/train/
  val_data: /path/to/category/val/
  num_classes: 1
  batch_size: 64
  workers: 4
  category: bike
  num_symmetry: 1
  max_objs: 10
        

To train a CenterPose model, use this command:

            

            
tao model centerpose 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 this argument is not specified, the results_dir from the spec file is used.

• -k, --key: A user-specific encoding key to save or load a .tlt model. If this argument is not specified, the model checkpoint isn’t encrypted.

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

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

Sample Usage

Here’s an example of the train command:

            

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

Optimizing Resource for Training CenterPose

Training CenterPose requires GPUs (for example, V100/A100) and CPU memory to be trained on a standard dataset, such as Objectron. The following are 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 trick is to reduce dataset.batch_size, which can cause your training to take longer than usual.

Typically, the following options result in a more balanced performance optimization:

• Set train.precision to fp16 to enable automatic mixed precision training. This can reduce your GPU memory usage and speed up the training. But might affect the accuracy.

• Try using more lightweight backbones like DLA34.

evaluate

The evaluate parameter defines the hyperparameters of the evaluate process.

            

            
evaluate:
  num_gpus: 1
  checkpoint: /path/to/model.pth
  opencv: False
  eval_num_symmetry: 1
  results_dir: /path/to/saving/directory
        

To run evaluation with a CenterPose model, use this command:

            

            
tao model centerpose 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 this value is not specified, a .pth model must be used.

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

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

Sample Usage

The following is an example of using the evaluate command:

            

            
tao model centerpose 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
  visualization_threshold: 0.3
  principle_point_x: 300.7
  principle_point_y: 392.8
  focal_length_x: 615.0
  focal_length_y: 615.0
  skew: 0.0
  use_pnp: True
  save_json: True
  save_visualization: True
  opencv: True
        

The inference tool for CenterPose models can be used to visualize 3D bounding boxes in 2D image plane, the order of points and the object relative dimension. Furthermore, it also generates a frame-by-frame JSON file for recording the results for each image.

            

            
tao model centerpose 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 this value is not specified, a .pth model must be used.

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

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

Sample Usage

The following is an example of using the inference command:

            

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

export

The export parameter defines the hyperparameters of the export process.

            

            
export:
  gpu_id: 0
  checkpoint: /path/to/model.pth
  onnx_file: /path/to/model.onnx
  input_channel: 3
  input_width: 512
  input_height: 512
  opset_version: 16
  do_constant_folding: True
        

            

            
tao model centerpose 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 this value is 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 is saved.

Sample Usage

The following is an example of using the export command:

            

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

For deployment, refer to TAO Deploy documentation for CenterPose.