tao-train-centerpose

CenterPose

独立安装？ 如果此会话并非由 TAO skill bank 插件初始化，请先运行 tao-setup 技能（主机预检、凭据、跨技能发现）。

用于关键点／姿态估计的 CenterPose。检测对象中心并回归关键点位置。用于 6-DoF 对象姿态估计。

设置 model.backbone.pretrained_backbone_path。

对于 TAO Deploy TensorRT 操作（gen_trt_engine、TensorRT evaluate 和 TensorRT inference），请使用此技能 references/ 文件夹中打包的部署规格模板，其前缀为 spec_template_deploy_*.yaml。

Dataclass 架构

生成的 TAO Core 架构打包在 schemas/<action>.schema.json 中，schemas/manifest.json 列出了可用操作。每个生成的架构还会根据架构顶层的 default 字段生成 references/spec_template_<action>.yaml。AutoML 启用状态通过 references/skill_info.yaml 模型层中的 automl_enabled 声明。某个操作要能运行 AutoML，必须存在并能解析 schemas/<action>.schema.json 和 references/spec_template_<action>.yaml。请使用打包的所选操作架构来获取 automl_default_parameters、automl_disabled_parameters、默认值、最小／最大边界、枚举、选项权重、数学条件、依赖项和常用参数。不要假设运行时存在 ~/tao-core；维护者会在打包 skill bank 前重新生成架构和模板。

训练操作策略

此模型在模型层启用了 AutoML。处理任何训练阶段的请求前，请先读取 references/skill_info.yaml，并根据明确指定的 automl_policy 值或用户的工作流请求确定本次运行的覆盖设置。默认使用 automl_policy: on，且在新的启动提示中只提供 on / off 选项。将“turn off AutoML”“disable AutoML”“no HPO”或“plain training”等表述视为仅对本次运行设置 automl_policy: off。当 automl_policy: on、automl_enabled: true，且已打包 schemas/train.schema.json 和 references/spec_template_train.yaml 时，默认通过 tao-skill-bank:tao-run-automl 执行训练操作，并使用此模型的 skill_dir。保留工作流／应用对数据集、规格、输出目录、GPU／平台设置、父检查点和 automl_policy 的覆盖设置。仅当 automl_policy: off 或缺少已打包的训练架构／模板时，才直接进行模型训练；如果缺少架构，请说明此模型已启用 AutoML，但在生成架构前无法运行 AutoML。

evaluate、inference、export 等非训练操作以及部署流程仍由此模型技能处理。每次运行的 automl_policy 覆盖设置不会更改模型元数据。

训练要求

• 数据集类型： centerpose
• 格式： 默认
• 训练监控指标： val_3DIoU、val_2DMPE
• 评估任务指标： test_3DIoU、test_2DMPE
• AutoML 指标约定： 对于基于评估的必需基线、通过 eval_fn 得出的每条建议、最佳模型选择以及最终评估，均使用 test_3DIoU 并采用最大化方向。训练会输出 val_3DIoU， 但评估操作不会输出，因此它仅适用于明确接受且不带影响基线的训练代理运行。

各操作的数据集要求

操作	规格键	来源	文件	列表？
evaluate	dataset.test_data	eval_dataset	test.tar.gz	否
gen_trt_engine	gen_trt_engine.tensorrt.calibration.cal_image_dir	calibration_dataset	train.tar.gz	是
inference	dataset.inference_data	inference_dataset	val.tar.gz	否
train	dataset.train_data	train_datasets	train.tar.gz	否
train	dataset.val_data	eval_dataset	val.tar.gz	否

典型规格覆盖项

每个操作都必须覆盖数据源——Agent 必须根据上方的 Per-Action Dataset Requirements 表构建数据源路径，并将其包含在 spec_overrides 中。

TRAIN_DIR = "/path/to/extracted/train"
VAL_DIR = "/path/to/extracted/val"
TEST_DIR = "/path/to/extracted/test"
INFER_DIR = VAL_DIR
CAL_IMAGE_DIRS = ["/path/to/extracted/train/<sequence_or_image_dir>"]

train（必需数据源）：

{
    "train.num_epochs": 30,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "train.num_gpus": 1,
    "dataset.category": "bike",
    "dataset.batch_size": 4,
    "dataset.train_data": TRAIN_DIR,
    "dataset.val_data": VAL_DIR,
}

evaluate（必需数据源）：

{
    "dataset.category": "bike",
    "dataset.test_data": TEST_DIR,
}

inference（必需数据源）：

{
    "dataset.category": "bike",
    "dataset.inference_data": INFER_DIR,
}

gen_trt_engine（必需数据源）：

{
    "gen_trt_engine.tensorrt.calibration.cal_image_dir": CAL_IMAGE_DIRS,
}

评估数据集

可选。验证集和测试集以单独的 tar 包提供。训练会写入 val_3DIoU/val_2DMPE KPI，而 evaluate 会写入 test_3DIoU/test_2DMPE；不要配置由 evaluate 驱动的 AutoML 回调 来提取带训练前缀的名称。

重要参数

• dataset.num_classes：对象类别数。默认值为 1。
• dataset.num_joints：每个对象的关键点数。固定为 8（bbox 关键点）。有效范围：只能为 8。
• dataset.input_res：输入分辨率。固定为 512。输出分辨率固定为 128。
• dataset.category：对象类别名称。默认值为 "cereal_box"。
• model.backbone.model_type：默认值为 fan_small。骨干网络选项受 schema 限制。
• train.optim.lr：学习率。默认值为 6e-5。使用 MultiStep 调度器，lr_steps=[90, 120]，lr_decay=0.1。
• train.loss_config：包含多个开关的丰富损失配置：mse_loss、obj_scale、obj_scale_uncertainty、hps_uncertainty、reg_bbox、hm_hp。权重：wh_weight=0.1、off_weight=1、hp_weight=1。
• inference.use_pnp：是否使用 PnP 进行 6-DoF 位姿估计。默认值为 True。需要相机内参（focal_length_x/y、principle_point_x/y）。
• export.input_width：导出输入尺寸。固定为 512x512。opset_version=16。

多 GPU / 多节点

**启动方式：**由 Lightning 管理（单个 python 进程，Lightning 会启动工作进程）。

Spec Key	描述	默认值
train.num_gpus	GPU 数量	1
train.gpu_ids	GPU 设备索引	[0]

• 策略：auto（Lightning 自动选择最佳策略）
• 没有显式的 num_nodes 或 distributed_strategy 配置——仅支持单节点
• 不支持 sync_batchnorm

导出 / TRT 默认值

• 导出输入：512x512（固定），opset 16
• TRT 数据类型：FP32、FP16、INT8
• TRT opt_batch_size：4，max_batch_size：8

硬件

至少需要 1 块 GPU，建议使用 2 块 GPU。每块 GPU 需有 16GB 以上显存。CenterPose 的显存占用适中，具体取决于输入分辨率和关键点数量。

错误模式

num_joints 不匹配：确保 dataset.num_joints 与标注中的关键点数一致。

为本地 Docker 解压 S3 tar 包：starter-kit 的 S3 数据打包为 train.tar.gz、val.tar.gz 和 test.tar.gz，但 CenterPose TAO 操作 需要解压后的文件夹。解压每个归档，并将 dataset.train_data、 dataset.val_data、dataset.test_data 和 dataset.inference_data 设置为 解压后的数据集拆分目录。

检查点交接：CenterPose 训练会写入具体的检查点，例如 model_epoch_000_step_00008.pth，以及一个 centerpose_model_latest.pth 符号链接。 对于 evaluate、inference、export 和 resume，请使用 SDK/模型检查点解析器或指定确切的轮次/步数检查点。仅当用户明确要求使用最新版本时，才使用该符号链接。

TAO Deploy 后处理器兼容性：使用从该技能固定的部署镜像或所选平台解析出的部署镜像。gen_trt_engine 成功运行 并不能证明部署环境中的 evaluate 或 inference 可正常运行；请分别检查这些操作的退出代码 和日志，尤其要留意 CenterPose 后处理器错误，例如 TypeError: only 0-dimensional arrays can be converted to Python scalars。

Spec 参数 / 父模型推断

特定于模型的推断映射应放在此 MD 文件中，而不是 config.json 中。生成的 runner 应读取本节，并在调用 create_job() 之前使用 SDK 辅助函数应用这些映射。这与旧版微服务的 infer_params.py 流程一致。

来自 TAO Core centerpose.config.json 的推断映射：

操作	Spec 字段	推断函数	含义
evaluate	encryption_key	key	加密密钥
evaluate	evaluate.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
evaluate	evaluate.trt_engine	parent_model	从父作业结果文件夹推断出的模型文件
evaluate	results_dir	output_dir	当前作业结果目录
export	encryption_key	key	加密密钥
export	export.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
export	export.onnx_file	create_onnx_file	输出 ONNX 路径
export	results_dir	output_dir	当前作业结果目录
gen_trt_engine	encryption_key	key	加密密钥
gen_trt_engine	gen_trt_engine.onnx_file	parent_model	从父作业结果文件夹推断出的模型文件
gen_trt_engine	gen_trt_engine.tensorrt.calibration.cal_cache_file	create_cal_cache	校准缓存路径
gen_trt_engine	gen_trt_engine.trt_engine	create_engine_file	输出 TensorRT 引擎路径
gen_trt_engine	results_dir	output_dir	当前作业结果目录
inference	encryption_key	key	加密密钥
inference	inference.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
inference	inference.trt_engine	parent_model	从父作业结果文件夹推断出的模型文件
inference	results_dir	output_dir	当前作业结果目录
train	encryption_key	key	加密密钥
train	model.backbone.pretrained_backbone_path	ptm_if_no_resume_model	不存在恢复检查点时使用的 PTM
train	results_dir	output_dir	当前作业结果目录
train	train.resume_training_checkpoint_path	resume_model	从当前作业结果文件夹推断出的模型文件

对于 parent_model 或 parent_model_folder，请将上游 train/export/AutoML 子作业 ID 作为 parent_job_id 传入。SDK 会列出父作业结果文件夹，筛选检查点制品，并返回选定的模型文件或文件夹。不要将这些映射加回 config.json，也不要修改生成的 runner 脚本来猜测检查点路径。

部署

• tao-deploy-centerpose