CenterPose

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

• train

• evaluate

• inference

• export

These tasks can be invoked from the TAO 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
  gpu_ids: [0]
  checkpoint_interval: 5
  validation_interval: 5
  num_epochs: 10
  clip_grad_val: 100.0
  seed: 1234
  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
  gpu_ids: [0]
  checkpoint_interval: 5
  validation_interval: 5
  num_epochs: 10
  clip_grad_val: 100.0
  seed: 1234
  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
        

Use the following command to run CenterPose training:

            

            
tao model centerpose train [-h] -e <experiment_spec_file>
                     [results_dir=<global_results_dir>]
                     [model.<model_option>=<model_option_value>]
                     [dataset.<dataset_option>=<dataset_option_value>]
                     [train.<train_option>=<train_option_value>]
                     [train.gpu_ids=<gpu indices>]
                     [train.num_gpus=<number of gpus>]
        

Required Arguments

The only required argument is the path to the experiment spec:

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

Optional Arguments

You can set optional arguments to override the option values in the experiment spec file.

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

• model.<model_option>: The model options.

• dataset.<dataset_option>: The dataset options.

• train.<train_option>: The train options.

• train.optim.<optim_option>: The optimizer options

Checkpointing and Resuming Training

At every train.checkpoint_interval, a PyTorch Lightning checkpoint is saved. It is called model_epoch_<epoch_num>.pth. These are saved in train.results_dir, like so:

            

            
$ ls /results/train

'model_epoch_000.pth'
'model_epoch_001.pth'
'model_epoch_002.pth'
'model_epoch_003.pth'
'model_epoch_004.pth'
        

The latest checkpoint will also be saved as centerpose_model_latest.pth. Training will automatically resume from centerpose_model_latest.pth if it exists in train.results_dir. This will be superseded by train.resume_training_checkpoint_path if it is provided.

The major implication of this logic is that, if you wish to trigger fresh training from scratch, either

• Specify a new, empty results directory (Recommended), or

• Remove the latest checkpoint from the results directory

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:
  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>
                     evaluate.checkpoint=<model to be evaluated>
                     [evaluate.<evaluate_option>=<evaluate_option_value>]
                     [evaluate.gpu_ids=<gpu indices>]
                     [evaluate.num_gpus=<number of gpus>]
        

Required Arguments

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

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

Optional Arguments

• evaluate.<evaluate_option>: The evaluate options.

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>
                     inference.checkpoint=<model to be inferenced>
                     [inference.<inference_option>=<inference_option_value>]
                     [inference.gpu_ids=<gpu indices>]
                     [inference.num_gpus=<number of gpus>]
        

Required Arguments

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

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

Optional Arguments

• inference.<inference_option>: The inference options.

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>
                     export.checkpoint=<model to export>
                     export.onnx_file=<onnx path>
                     [export.<export_option>=<export_option_value>]
        

Required Arguments

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

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

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

Optional Arguments

• export.<export_option>: The export options.

For deployment, refer to TAO Deploy documentation for CenterPose.