tao-train-pose-classification

姿态分类

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

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

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

打包的 PyTorch Pose Classification 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。

快速开始（docker run）

Docker 原生启动——主机上无需 TAO SDK，也无需 Python。当本地 Docker/平台技能提供了更严格的环境专用命令（非 root UID 映射、缓存重定向、远程守护进程）时，请改用该技能。

TAO_PYT_IMAGE_DEFAULT=nvcr.io/nvidia/tao/tao-toolkit:7.3.0-pyt  # versions-key: images.tao_toolkit.pyt
TAO_PYT_IMAGE="${TAO_PYT_IMAGE:-$TAO_PYT_IMAGE_DEFAULT}"
RUN_ROOT="${RUN_ROOT:-$PWD}"
DOCKER_COMMON=(
  --rm --gpus all --shm-size=8g
  --shm-size=8g
  --ulimit memlock=-1
  --ulimit stack=67108864
  -v "$RUN_ROOT/data:/data:ro"
  -v "$RUN_ROOT/specs:/specs:ro"
  -v "$RUN_ROOT/results:/results"
)

数据集转换：

docker run "${DOCKER_COMMON[@]}" "$TAO_PYT_IMAGE" \
  pose_classification dataset_convert -e /specs/dataset_convert.yaml

训练：

docker run "${DOCKER_COMMON[@]}" "$TAO_PYT_IMAGE" \
  pose_classification train -e /specs/train.yaml

评估：

docker run "${DOCKER_COMMON[@]}" "$TAO_PYT_IMAGE" \
  pose_classification evaluate -e /specs/evaluate.yaml

推理：

docker run "${DOCKER_COMMON[@]}" "$TAO_PYT_IMAGE" \
  pose_classification inference -e /specs/inference.yaml

导出：

docker run "${DOCKER_COMMON[@]}" "$TAO_PYT_IMAGE" \
  pose_classification export -e /specs/export.yaml

每项操作都通过 -e 接收其 spec；results_dir 在 spec 中设置，也可在命令行中覆盖。挂载 spec 引用的任何预训练权重目录，并确保各项操作中容器内路径保持一致。

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 选项。将“关闭 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。保留工作流/应用层对数据集、spec、输出目录、GPU/平台设置、父检查点以及 automl_policy 的覆盖。仅当 automl_policy: off 或缺少已打包的训练架构/模板时，才直接进行模型训练；若缺少架构/模板，请报告 AutoML 已启用，但在为此模型生成架构之前无法运行。

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

训练要求

• 数据集类型： pose_classification
• 格式： default
• AutoML 训练指标： val_loss，其中 direction=minimize
• 独立评估指标： accuracy，其中 direction=maximize。使用 val_loss 对训练阶段的 AutoML 推荐结果排序，仅使用 accuracy 验证所选检查点能否成功加载和评估。

每个操作的数据集要求

操作	Spec Key	来源	文件	列表？
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",
}

数据集转换

Pose Classification 的数据集转换是可选的。只有当用户提供原始 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；在训练/评估/推理/导出流程中直接使用这些文件，不要合成虚假的 BodyPose JSON。

评估数据集

可选。训练时会同时提供验证数据：val_data.npy / val_label.pkl。对于此模型，TAO 训练会将 val_loss 作为 TensorBoard 验证标量输出；除非自定义评估钩子提供了其他指标，否则 AutoML 选择时使用方向为 minimize 的 val_loss。独立 evaluate 操作会输出 accuracy；请将该值与所选检查点的评估结果进行比较，但不要用它替代训练阶段用于 KPI 排序的 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，也不要修改生成的运行器脚本来猜测检查点路径。