⚡

姿态分类训练（PoseClassification）Skill tao-train-pose-classification

该技能提供基于 ST-GCN（时空图卷积网络）的姿态分类全流程：将骨架关键点序列数据集转换为 TAO 格式（可选）、训练分类模型、评估、导出 ONNX 和推理。包含 AutoML 调参、多 GPU 支持与参数覆盖指南。关键词：姿态分类、ST-GCN、骨架序列识别、关键点分类、TAO、动作识别、val_loss。

视觉模型训练 0 次安装 9 次浏览 更新于 9/6/2026
名称 tao-train-pose-classification
描述 姿态分类（Pose Classification）。使用 ST-GCN（时空图卷积网络）对骨架序列进行分类，将姿态关键点数据映射为动作类别。
开源协议 Apache-2.0

姿态分类

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

使用 ST-GCN（Spatial Temporal Graph Convolutional Network）进行姿态分类：将骨架序列分类为动作类别。可从关键点数据从头训练。

打包的 PyTorch 姿态分类 CLI 支持 dataset_convert、train、evaluate、export 和 inference。dataset_convert 仅在输入为原始 DeepStream BodyPose JSON 时执行。若数据集已转换为 TAO 就绪的 .npy / .pkl 文件，则直接使用这些文件进行 train，并在验证报告中注明“未运行数据集转换：数据集已预转换”。不提供 deploy、prune、quantize 或独立 retrain。Resume/retrain 通过 pose_classification train -e ... 并设置 train.resume_training_checkpoint_path 实现。

快速开始（docker run）

TAO_PYT_IMAGE_DEFAULT=nvcr.io/nvidia/tao/tao-toolkit:7.1.0-pyt
TAO_PYT_IMAGE=${TAO_PYT_IMAGE:-$TAO_PYT_IMAGE_DEFAULT}
RUN_ROOT=${RUN_ROOT:-$PWD}
DOCKER_COMMON=(
  --rm --gpus all --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

评估、推理、导出类似。

Dataclass Schemas

生成 TAO Core schemas 打包在 schemas/<action>.schema.json 中，schemas/manifest.json 列出可用操作。每个 schema 也通过顶层 default 导出 references/spec_template_<action>.yaml。AutoML 在 references/skill_info.yaml 中以模型层声明；automl_enabled 为 true。只有同时存在 schema 和 spec template 且能解析时，AutoML 才可运行。

Train Action Policy

该模型在模型层启用 AutoML。处理训练请求前，读取 references/skill_info.yaml，从明确值 automl_policy 或用户工作流中解析运行覆盖。默认使用 automl_policy: on，并仅在启动提示中公开 on / 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/platform 设置、父 checkpoint 和 automl_policy 的覆盖。若 automl_policy: off 或缺少对应 schema/template，则走直接模型训练；若缺少 schema，则报告 AutoML 已启用但暂不可运行。

非训练动作（evaluate/inference/export/deploy）仍留在本模型技能中。单次运行的覆盖不会改变模型元数据。

训练要求

数据集类型：pose_classification；格式：默认；监控指标：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_overrides 中。若已提供 .npy/.pkl，则不要运行 dataset_convert。

常用 Spec 覆盖

数据源覆盖是强制性的。例如：

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

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_5: 5},
    model.graph_layout: nvidia,
    dataset.train_dataset.data_path: {S3_TRAIN}/train_data.npy,
    dataset.train_dataset.label_path: {S3_TRAIN}/train_label.pkl,
    dataset.val_dataset.data_path: {S3_TRAIN}/val_data.npy,
    dataset.val_dataset.label_path: {S3_TRAIN}/val_label.pkl,
}

Resume、evaluate、export、inference 的 spec 也按要求设置 checkpoint 和输出路径。

数据集转换

转换是可选的，仅在用户提供原始 DeepStream BodyPose JSON 时运行。常见 S3 验证数据已转换，直接使用 .npy/.pkl，不要合成伪造 JSON。

Eval 数据集

可选。验证数据随训练提供为 val_data.npy / val_label.pkl。TAO 训练会输出 val_loss 作为 TensorBoard 验证标量，AutoML 选择最小化 val_loss 的 trial。

重要参数

  • dataset.num_classes: 动作类别数，默认 6。
  • model.graph_layout: 骨架图布局。选项：nvidia、openpose。决定关节点连接。
  • model.graph_strategy: 用于 GCN 的图划分策略。
  • train.optim.lr: 学习率。默认 0.1（SGD）。
  • model.dropout: Dropout 比例。

多 GPU/多节点

Lightning 管理单进程；单机多卡。

Spec Key Description Default
train.num_gpus GPU 数量 1
train.gpu_ids GPU device indices [0]

轻量模型，通常单卡足够。

硬件

最少 1 张 GPU；推荐 1 张；8GB+ 显存。姿态分类数据小，轻量。

错误模式

  • 图布局不匹配：确保 model.graph_layout 与 .npy 文件中骨架格式匹配。
  • 标签形状不匹配：class 索引必须在 [0, num_classes)。
  • 缺少 label map：训练 dataloader 期望 dataset.label_map 为字典。若只有数值类 ID，请构造合成连续映射，如 class_0: 0 至 class_5: 5。
  • 检查点交接：从父结果文件夹中选择所需 .pth checkpoint，例如 model_epoch_000_step_00007.pth，并传给 evaluate/export/inference 或 resume 路径。pc_model_latest.pth 是 latest 符号链接，仅在用户明确要求最新时使用。下游保持相同 dataset.num_classes、dataset.label_map、model.graph_layout。
  • 动作特定路径：evaluate 和 inference 模板中仍包含训练 dataset 块。evaluate 要设置 evaluate.test_dataset.data_path 和 label_path；inference 要设置 inference.test_dataset.data_path 和 inference.output_file，不要只改第一个。

Spec Param / Parent Model Inference

推理映射表中的字段（如 evaluate.checkpoint -> parent_model、export.checkpoint -> parent_model、inference.output_file 等）通过父任务结果目录推断。应传上游 train/export/AutoML 子任务 id 作为 parent_job_id。SDK 列出父结果文件夹，筛选 checkpoint 工件，返回模型文件或文件夹。不要将这些映射写回 config.json。