CenterPose
CenterPose 是 TAO 中包含的类别级物体姿态估计模型。它支持以下任务
训练
评估
推理
导出
这些任务可以使用 TAO Launcher 从命令行调用，使用以下约定
tao model centerpose <sub_task> <args_per_subtask>
其中，args_per_subtask 是给定子任务所需的命令行参数。以下章节将详细解释每个子任务。
CenterPose 期望图像目录和带注释的 JSON 文件用于训练或验证。有关输入数据格式的更多信息，请参阅 CenterPose 数据格式 页面。
CenterPose 的训练实验规范文件包括 model、train 和 dataset 参数。这是一个使用 fan_small 骨干网络在 Google Objectron 数据集 bike 类别上训练 CenterPose 模型的示例规范文件。
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 |
dict config | – | 模型架构的配置 | |
dataset |
dict config | – | 数据集的配置 | |
训练 |
dict config | – | 训练任务的配置 | |
评估 |
dict config | – | 评估任务的配置 | |
推理 |
dict config | – | 推理任务的配置 | |
encryption_key |
string | None | 用于加密和解密模型文件的加密密钥 | |
results_dir |
string | /results | 实验结果保存的目录 | |
导出 |
dict config | – | ONNX 导出任务的配置 | |
gen_trt_engine` |
dict config | – | TensorRT 生成任务的配置。仅在 tao deploy 中使用 |
model
model 参数提供了更改 CenterPose 架构的选项。
model:
down_ratio: 4
use_pretrained: False
backbone:
model_type: fan_small
pretrained_backbone_path: /path/to/your-fan-small-pretrained-model
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
down_ratio |
int | 4 | 网络特征图的下采样比例。 | 4 |
use_pretrained |
bool | False | 一个标志，指定是否使用预训练权重初始化骨干网络。 | True, False |
backbone |
dict config | 骨干网络模型类型的配置和预训练权重的路径。 | >0 |
backbone
backbone 参数提供了更改 CenterPose 骨干网络架构的选项。
backbone:
model_type: fan_small
pretrained_backbone_path: /path/to/your-fan-small-pretrained-model
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
pretrained_backbone_path |
string | None | 预训练骨干网络文件的可选路径。使用 “FAN” 骨干网络时设置预训练路径。“DLA34” 骨干网络可以自动下载预训练权重，将其设置为 “null”。 | string to the path |
model_type |
string | DLA34 | 模型的骨干网络名称。支持 DLA34 和 FAN。 | DLA34, fan_small, fan_base, fan_large |
训练
train 参数定义了训练过程的超参数。
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]
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
num_gpus |
unsigned int | 1 | 用于分布式训练的 GPU 数量 | >0 |
gpu_ids |
List[int] | [0] | 用于分布式训练的 GPU 索引 | |
seed |
unsigned int | 1234 | 用于 random、numpy 和 torch 的随机种子 | >0 |
num_epochs |
unsigned int | 10 | 运行实验的总 epoch 数 | >0 |
checkpoint_interval |
unsigned int | 1 | 保存检查点的 epoch 间隔 | >0 |
validation_interval |
unsigned int | 1 | 运行验证的 epoch 间隔 | >0 |
resume_training_checkpoint_path |
string | 要从中恢复训练的中间 PyTorch Lightning 检查点 | ||
results_dir |
string | /results/train | 保存训练结果的目录 | |
clip_grad_val |
float | 100.0 | 在指定值处裁剪参数迭代器的梯度 | >=0 |
precision |
string | fp32 | 指定 “fp16” 启用混合精度训练。使用 fp16 训练可以帮助节省 GPU 内存。 | fp32, fp16 |
optim |
dict config | 优化器的配置，包括学习率、学习率调度器 | >0 |
optim
optim 参数定义了训练中优化器的配置，包括学习率和学习率步长。
optim:
lr: 6e-05
lr_steps: [90, 120]
参数 |
数据类型 |
默认值 |
描述 |
支持的值 |
|---|---|---|---|---|
lr |
float | 6e-05 | 训练模型的初始学习率，不包括骨干网络 | >0.0 |
lr_steps |
int list | [90, 120] | 用于调度器降低学习率的步长 | int list |
dataset
dataset 参数定义了数据集来源、训练批大小和数据集设置。
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_data |
string | 训练数据路径：包含训练图像及其相关 JSON 文件的目录。图像和 JSON 文件在同一文件夹中使用相同的文件名 | ||
val_data |
string | 验证数据路径：包含验证图像及其相关 JSON 文件的目录。图像和 JSON 文件在同一文件夹中使用相同的文件名 | ||
test_data |
string | 测试数据路径：包含测试图像及其相关 JSON 文件的目录。图像和 JSON 文件在同一文件夹中使用相同的文件名 | ||
inference data |
string | 推理数据路径：包含推理图像的目录。推理管道不需要 JSON 文件 | ||
num_classes |
unsigned int | 1 | 训练数据中的类别数量。由于 CenterPose 是一种类别级姿态估计方法，因此仅支持 1 个类别。 | 1 |
batch_size |
unsigned int | 4 | 训练和验证的批大小 | >0 |
workers |
unsigned int | 8 | 并行处理数据的工作进程数 | >0 |
category |
string | 训练数据集的类别名称。不同类别可能有不同的训练策略。请参阅 num_symmetry 了解更多详情 |
||
num_symmetry |
unsigned int | 1 | 对称旋转的次数，这意味着 3D 边界框沿 y 轴旋转的次数。每个旋转的边界框都被视为训练的真值。例如，bottle 是对称物体，num_symmetry 可以设置为 12（每次旋转 30 度）。当物体非对称时，num_symmetry 设置为 1 |
>0 |
max_objs |
unsigned int | 10 | 单张图像中用于训练的最大物体数量。 | >0 |
使用以下命令运行 CenterPose 训练
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>]
必需参数
唯一必需的参数是实验规范的路径
-e, --experiment_spec: 用于设置训练实验的实验规范文件
可选参数
您可以设置可选参数来覆盖实验规范文件中的选项值。
-h, --help: 显示此帮助消息并退出。model.<model_option>: 模型选项。dataset.<dataset_option>: 数据集选项。train.<train_option>: 训练选项。train.optim.<optim_option>: 优化器选项
对于训练、评估和推理，我们为每个相应的任务公开 2 个变量：num_gpus 和 gpu_ids，它们分别默认为 1 和 [0]。如果两者都已传递，但彼此不一致，例如 num_gpus = 1，gpu_ids = [0, 1]，则它们将被修改为遵循具有更多 GPU 的设置，例如 num_gpus = 1 -> num_gpus = 2。
检查点和恢复训练
在每个 train.checkpoint_interval，都会保存一个 PyTorch Lightning 检查点。它被称为 model_epoch_<epoch_num>.pth。这些检查点保存在 train.results_dir 中，如下所示
$ ls /results/train
'model_epoch_000.pth'
'model_epoch_001.pth'
'model_epoch_002.pth'
'model_epoch_003.pth'
'model_epoch_004.pth'
最新的检查点也将保存为 centerpose_model_latest.pth。如果 centerpose_model_latest.pth 存在于 train.results_dir 中，训练将自动从它恢复。如果提供了 train.resume_training_checkpoint_path，则它将被 train.resume_training_checkpoint_path 取代。
此逻辑的主要含义是，如果您希望从头开始触发全新训练，则可以
指定一个新的、空的 results 目录（推荐），或者
从 results 目录中删除最新的检查点
优化 CenterPose 训练的资源
训练 CenterPose 需要 GPU（例如，V100/A100）和 CPU 内存才能在标准数据集（例如 Objectron）上进行训练。以下是一些您可以使用的策略，以便仅使用有限的资源启动训练。
优化 GPU 内存
有多种方法可以优化 GPU 内存使用量。一种技巧是减少 dataset.batch_size，这可能会导致您的训练比平常花费更长的时间。
通常，以下选项会带来更均衡的性能优化
将
train.precision设置为fp16以启用自动混合精度训练。这可以减少您的 GPU 内存使用量并加快训练速度。但可能会影响准确性。尝试使用更轻量级的骨干网络，如
DLA34。
评估
evaluate 参数定义了评估过程的超参数。
evaluate:
checkpoint: /path/to/model.pth
opencv: False
eval_num_symmetry: 1
results_dir: /path/to/saving/directory
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
checkpoint |
string | 要评估的 PyTorch 模型路径 | ||
results_dir |
string | /results/evaluate | 保存评估结果的目录 | |
num_gpus |
unsigned int | 1 | 用于分布式评估的 GPU 数量 | >0 |
gpu_ids |
List[int] | [0] | 用于分布式评估的 GPU 索引 | |
|
|
bool |
False |
如果 |
True, False |
|
|
unsigned int |
1 |
对于对称物体类别（例如 bottle），我们沿对称轴将估计的边界框旋转 N 次（N = 100），并针对每个旋转的实例评估预测结果 |
>0 |
trt_engine |
string | 要评估的 TensorRT 模型路径。应仅与 tao deploy 一起使用 |
要使用 CenterPose 模型运行评估，请使用此命令
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>]
必需参数
-e, --experiment_spec: 用于设置评估实验的实验规范文件evaluate.checkpoint: 要评估的.pth模型。
可选参数
evaluate.<evaluate_option>: 评估选项。
推理
inference 参数定义了推理过程的超参数。
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
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
checkpoint |
string | 要推理的 PyTorch 模型路径 | ||
results_dir |
string | /results/inference | 保存推理结果的目录 | |
num_gpus |
unsigned int | 1 | 用于分布式推理的 GPU 数量 | >0 |
gpu_ids |
List[int] | [0] | 用于分布式推理的 GPU 索引 | |
visualization_threshold |
float | 0.3 | 用于过滤预测结果的置信度阈值 | >=0 |
principle_point_x |
float | 300.7 | 内参矩阵的主点 x 坐标。请使用与您的数据相匹配的正确相机标定矩阵 | >0 |
principle_point_y |
float | 392.8 | 内参矩阵的主点 y 坐标。请使用与您的数据相匹配的正确相机标定矩阵 | >0 |
focal_length_x |
float | 615.0 | 内参矩阵的焦距 x。请使用与您的数据相匹配的正确相机标定矩阵 | >0 |
focal_length_y |
float | 615.0 | 内参矩阵的焦距 y。请使用与您的数据相匹配的正确相机标定矩阵 | >0 |
skew |
float | 0.0 | 内参矩阵的倾斜系数。请使用与您的数据相匹配的正确相机标定矩阵 | >=0 |
use_pnp |
bool | True | 用于建立 2D-3D 对应关系以求解 6-DoF 姿态的 PnP 算法 | True, False |
save_json |
bool | True | 将所有结果保存到本地 JSON 文件，包括 2D 关键点、3D 关键点、位置、四元数 和 相对尺度 | True, False |
|
|
bool |
True |
将可视化结果保存到本地 .jpg 文件，包括 投影的 2D 边界框 以及点顺序、相对尺度 和 物体姿态 |
True, False |
|
|
bool |
False |
如果 |
True, False |
trt_engine |
string | 要推理的 TensorRT 模型路径。应仅与 tao deploy 一起使用 |
CenterPose 模型的推理工具可用于可视化 2D 图像平面中的 3D 边界框、点的顺序和物体相对尺寸。此外，它还会生成逐帧 JSON 文件，用于记录每张图像的结果。
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>]
必需参数
-e, --experiment_spec: 用于设置推理实验的实验规范文件inference.checkpoint: 要推理的.pth模型。
可选参数
inference.<inference_option>: 推理选项。
导出
export 参数定义了导出过程的超参数。
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
| 参数 | 数据类型 | 默认值 | 描述 | 支持的值 |
gpu_id |
unsigned int | 0 | 用于将 pth 模型转换为 ONNX 模型的 GPU ID | >=0 |
checkpoint |
string | 要导出的 PyTorch 模型路径 | ||
onnx_file |
string | .onnx 文件的路径 |
||
input_channel |
unsigned int | 3 | 输入通道大小。仅支持值 3。 | 3 |
input_width |
unsigned int | 512 | 输入宽度 | >0 |
input_height |
unsigned int | 512 | 输入高度 | >0 |
opset_version |
unsigned int | 16 | 导出的 ONNX 的 opset 版本 | >0 |
do_constant_folding |
bool | True | 是否执行常量折叠。如果 TensorRT 版本低于 8.6，则设置为 True | True, False |
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>]
必需参数
-e, --experiment_spec: 实验规范文件的路径export.checkpoint: 要导出的.pth模型。export.onnx_file: 保存.etlt或.onnx模型的路径。
可选参数
export.<export_option>: 导出选项。