tao-train-pose-classification

姿态分类

使用 ST-GCN（空间时间图卷积网络）进行姿态分类。根据姿态关键点数据，将骨架序列分类到不同的动作类别。

通常使用骨架数据从头开始训练。

打包的 PyTorch 姿态分类 CLI 支持 dataset_convert、train、evaluate、export 和 inference。dataset_convert 为条件操作：仅当输入是原始 DeepStream BodyPose JSON 时才运行。如果数据集已经转换为适用于 TAO 的 .npy / .pkl 文件，则直接对这些文件运行 train，并在验证报告中将数据集转换标记为 not run: preconverted dataset provided。此模型不提供 deploy、prune、quantize 或独立 retrain 操作。恢复训练/重新训练时，使用 pose_classification train -e ...，并填写 train.resume_training_checkpoint_path。

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/train.schema.json 和 references/spec_template_train.yaml。automl_default_parameters、automl_disabled_parameters、默认值、最小/最大边界、枚举、选项权重、数学条件、依赖项和常用参数均以打包的训练架构为准。运行时不要依赖 ~/tao-core；维护者应在打包技能库之前重新生成架构/模板。

训练操作策略

此模型在模型层启用了 AutoML。处理任何训练阶段的请求前，请读取 references/skill_info.yaml，并根据明确指定的 automl_policy 值或用户的工作流请求确定本次运行的策略。默认使用 automl_policy: on，新建启动提示时仅提供 on / off 两个选项。将“关闭 AutoML”“禁用 AutoML”“不使用 HPO”或“普通训练”等表述视为仅对本次运行设置 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 等非训练操作以及 deploy 流程仍在此模型技能中执行。每次运行的 automl_policy 覆盖设置不会更改模型元数据。

训练要求

• 数据集类型： pose_classification
• 格式： default
• 监控指标： val_loss

各操作的数据集要求

操作	Spec 键	来源	文件	列表？
dataset_convert（可选）	dataset_convert.data	id	DeepStream BodyPose JSON	否
evaluate	evaluate.test_dataset.data_path	train_datasets	val_data.npy	否
evaluate	evaluate.test_dataset.label_path	train_datasets	val_label.pkl	否
inference	inference.test_dataset.data_path	train_datasets	test_data.npy	否
train	dataset.train_dataset.data_path	train_datasets	train_data.npy	否
train	dataset.train_dataset.label_path	train_datasets	train_label.pkl	否
train	dataset.val_dataset.data_path	train_datasets	val_data.npy	否
train	dataset.val_dataset.label_path	train_datasets	val_label.pkl	否

典型的 Spec 覆盖项

每次运行操作都必须覆盖数据源——Agent 必须根据上方的「各操作的数据集要求」表构建数据源路径，并将其包含在 spec_overrides 中。如果提供的数据集已经转换为 .npy / .pkl 文件，则不要运行 dataset_convert。

S3_TRAIN = "s3://bucket/data/purpose_built_models_pose_classification_train/nvidia"
CHECKPOINT = "/results/{train_job_id}/results_dir/model_epoch_000_step_00007.pth"

dataset_convert（可选；仅适用于原始 DeepStream BodyPose JSON）：

{
    "dataset_convert.data": "s3://bucket/data/<deepstream-bodypose-output>.json",
}

train（必需的数据源）：

{
    "train.num_epochs": 30,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "train.num_gpus": 1,
    "wandb.enable": False,
    "dataset.num_classes": 6,
    "dataset.label_map": {
        "class_0": 0,
        "class_1": 1,
        "class_2": 2,
        "class_3": 3,
        "class_4": 4,
        "class_5": 5,
    },
    "model.graph_layout": "nvidia",
    "dataset.train_dataset.data_path": f"{S3_TRAIN}/train_data.npy",
    "dataset.train_dataset.label_path": f"{S3_TRAIN}/train_label.pkl",
    "dataset.val_dataset.data_path": f"{S3_TRAIN}/val_data.npy",
    "dataset.val_dataset.label_path": f"{S3_TRAIN}/val_label.pkl",
}

resume train（必需的检查点）：

{
    "train.num_epochs": 31,
    "train.resume_training_checkpoint_path": CHECKPOINT,
    "dataset.train_dataset.data_path": f"{S3_TRAIN}/train_data.npy",
    "dataset.train_dataset.label_path": f"{S3_TRAIN}/train_label.pkl",
    "dataset.val_dataset.data_path": f"{S3_TRAIN}/val_data.npy",
    "dataset.val_dataset.label_path": f"{S3_TRAIN}/val_label.pkl",
}

evaluate（必需的数据源）：

{
    "evaluate.test_dataset.data_path": f"{S3_TRAIN}/val_data.npy",
    "evaluate.test_dataset.label_path": f"{S3_TRAIN}/val_label.pkl",
    "evaluate.checkpoint": CHECKPOINT,
}

export（必需的检查点和输出）：

{
    "export.checkpoint": CHECKPOINT,
    "export.onnx_file": "/results/{export_job_id}/results_dir/pose_classification.onnx",
}

inference（必需的数据源）：

{
    "inference.test_dataset.data_path": f"{S3_TRAIN}/test_data.npy",
    "inference.test_dataset.label_path": f"{S3_TRAIN}/test_label.pkl",
    "inference.checkpoint": CHECKPOINT,
    "inference.output_file": "/results/pose_classification_inference.txt",
}

数据集转换

姿态分类的数据集转换是可选的。仅当用户提供原始 DeepStream BodyPose JSON 时，才运行 pose_classification dataset_convert。对于常见的 S3 验证数据集，数据已经转换为 train_data.npy、train_label.pkl、val_data.npy、val_label.pkl、test_data.npy 和 test_label.pkl；在 train/evaluate/inference/export 流程中直接使用这些文件，不要伪造 BodyPose JSON。

评估数据集

可选。验证数据与训练数据一同提供，文件为 val_data.npy / val_label.pkl。对于此模型，TAO 训练会将 val_loss 作为 TensorBoard 验证标量输出；除非自定义评估钩子提供了不同的指标，否则在 AutoML 选择中使用 val_loss，并将方向设为最小化。

重要参数

• dataset.num_classes：姿态动作类别的数量。默认值为 6。
• model.graph_layout：骨架图布局。可选项：nvidia、openpose。决定关节的连接方式。
• model.graph_strategy：GCN 的图划分策略。
• train.optim.lr：学习率。默认值为 0.1（SGD）。由于图卷积的特性，该值高于视觉模型。
• model.dropout：用于正则化的 Dropout 比率。

多 GPU / 多节点

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

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

• 策略：auto（Lightning 会自动选择最佳策略）
• 没有显式的 num_nodes 或 distributed_strategy 配置——仅支持单节点
• 模型轻量，通常使用单个 GPU 即可

硬件

至少需要 1 个 GPU，建议使用 1 个 GPU。每个 GPU 需配备 8GB 或更大的显存。姿态分类非常轻量——骨架数据很小。单个 GPU 即可满足需求。

错误模式

图布局不匹配：确保 model.graph_layout 与 .npy 数据文件中的骨架格式一致。

标签形状不匹配：train_label.pkl 的类别索引必须在 [0, num_classes) 范围内。

缺少标签映射：训练数据加载器要求 dataset.label_map 是一个字典。如果数据集只提供数字类别 ID，请设置一个连续的合成映射，例如，为六分类 NVIDIA 示例数据设置 class_0: 0 到 class_5: 5。

检查点交接：完成 AutoML/train 后，使用检查点解析器从父结果文件夹中选择预期保存的 .pth 检查点，例如 model_epoch_000_step_00007.pth，并将该确切文件作为 evaluate.checkpoint、export.checkpoint、inference.checkpoint 或 train.resume_training_checkpoint_path 传入。pc_model_latest.pth 是指向最新检查点的符号链接；只有当用户明确要求使用最新检查点，而不是指定或最佳检查点时，才使用它。后续操作中保持相同的 dataset.num_classes、dataset.label_map 和 model.graph_layout 覆盖项。

数据集转换来源：dataset_convert 需要 DeepStream BodyPose 应用输出的原始 JSON。常见的 NVIDIA 示例 S3 文件夹已转换为 train_data.npy、train_label.pkl、val_data.npy、val_label.pkl、test_data.npy 和 test_label.pkl；如果这些文件已存在，请跳过转换，直接使用转换后的文件。

特定操作的数据集路径：evaluate 和 inference 模板还包含训练用的 dataset.train_dataset 和 dataset.val_dataset 块。对于 evaluate，请填写 evaluate.test_dataset.data_path 和 evaluate.test_dataset.label_path。对于 inference，请填写 inference.test_dataset.data_path 并设置 inference.output_file；不要在替换文件中的第一个 data_path 或 label_path 后就停止。

输出文件：导出需要显式设置 export.onnx_file 路径。推理必须将 inference.output_file 设置为可写入的文件路径；打包模板的默认值为空字符串，当前 PyTorch 推理代码会直接打开该值。

规格参数／父模型推断

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

TAO Core pose_classification.config.json 中的推断映射：

操作	规格字段	推断函数	含义
dataset_convert	dataset_convert.results_dir	output_dir	当前作业的结果目录
evaluate	encryption_key	key	加密密钥
evaluate	evaluate.checkpoint	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	当前作业的结果目录
inference	encryption_key	key	加密密钥
inference	inference.checkpoint	parent_model	根据父作业结果文件夹推断出的模型文件
inference	inference.output_file	create_inference_result_file_pose	姿态推断结果文件
inference	results_dir	output_dir	当前作业的结果目录
train	encryption_key	key	加密密钥
train	model.pretrained_model_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，也不要修改生成的运行器脚本来猜测检查点路径。