mcpField
T
tao-train-reid

tao-train-reid

github.com/promptingcompany数据与数据库039 次安装17 周前更新
代码仓库

关于 tao-train-reid

用于行人重识别的 TAO 技能,通过学习判别性嵌入,在不同摄像头视角下匹配同一人。支持 AutoML 训练流程及基于检索的评估。

主要功能

  1. 01跨摄像头行人匹配
  2. 02判别性嵌入学习
  3. 03支持 AutoML 的训练流程
  4. 04CMC 与 mAP 评估

适用场景

  1. 01跨摄像头查找同一行人
  2. 02用查询集和图库评估身份匹配

安装

复制下面的命令,在终端中运行。

$npx skills add https://github.com/promptingcompany/nv-skills --skill tao-train-reid

安装后在 Claude Code 中输入/tao-train-reid确认已加载。

包含的文件

除 SKILL.md 外,安装时一起下载的参考资料和脚本。

BENCHMARK.md
skill-card.md
skill.oms.sig
evals/
evals.json
references/
skill_info.yaml
spec_template_evaluate.yaml
spec_template_export.yaml
spec_template_inference.yaml
spec_template_train.yaml
schemas/
evaluate.schema.json
export.schema.json
inference.schema.json
manifest.json
train.schema.json

详细介绍

再识别

行人再识别。学习具有判别性的嵌入,以匹配不同摄像头视角下的同一行人。基于度量学习。

为预训练权重设置 model.pretrained_model_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。使用打包的 train 架构获取 automl_default_parameters、automl_disabled_parameters、默认值、最小值/最大值边界、枚举值、选项权重、数学条件、依赖项和常用参数。不要期待运行时存在 ~/tao-core;维护者会在打包技能库前重新生成架构/模板。

训练操作策略

此模型在模型层级启用了 AutoML。在处理任何训练阶段请求前,读取 references/skill_info.yaml,并根据显式的 automl_policy 值或用户的工作流请求确定本次运行的覆盖设置。将“turn off AutoML”“disable AutoML”“no HPO”或“plain training”等表述视为仅对本次运行设置 automl_policy: off;否则默认为 auto。当 automl_policy: auto、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 已启用,但在生成此模型的架构前无法运行。

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

训练要求

  • 数据集类型: re_identification
  • 格式: default
  • 监控指标: cmc_rank_1
各操作的数据集要求
操作规格键来源文件列表?
evaluateevaluate.test_datasettrain_datasetssample_test.tar.gz否
evaluateevaluate.query_datasettrain_datasetssample_query.tar.gz否
inferenceinference.test_datasettrain_datasetssample_test.tar.gz否
inferenceinference.query_datasettrain_datasetssample_query.tar.gz否
traindataset.train_dataset_dirtrain_datasetssample_train.tar.gz否
traindataset.test_dataset_dirtrain_datasetssample_test.tar.gz否
traindataset.query_dataset_dirtrain_datasetssample_query.tar.gz否
典型规格覆盖

每项操作都必须覆盖数据源——Agent 必须根据上方的各操作数据集要求表构造数据源路径,并将其加入 spec_overrides。

S3_TRAIN = "s3://bucket/data/train"

train(必需数据源):

{
    "train.num_epochs": 30,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "train.num_gpus": 1,
    "num_classes": 100,
    "num_workers": 4,
    "batch_size": 16,
    "dataset.train_dataset_dir": f"{S3_TRAIN}/sample_train.tar.gz",
    "dataset.test_dataset_dir": f"{S3_TRAIN}/sample_test.tar.gz",
    "dataset.query_dataset_dir": f"{S3_TRAIN}/sample_query.tar.gz",
}

evaluate(必需数据源):

{
    "evaluate.test_dataset": f"{S3_TRAIN}/sample_test.tar.gz",
    "evaluate.query_dataset": f"{S3_TRAIN}/sample_query.tar.gz",
}

inference(必需数据源):

{
    "inference.test_dataset": f"{S3_TRAIN}/sample_test.tar.gz",
    "inference.query_dataset": f"{S3_TRAIN}/sample_query.tar.gz",
}

评估数据集

必需。基于检索的指标(CMC、mAP)评估需要测试数据集和查询数据集。

重要参数

  • dataset.num_classes:身份数量。默认值为 751。必须与训练数据中的唯一身份数量一致。
  • model.backbone:默认值为 resnet_50。
  • optim.base_lr:基础学习率。默认值为 3.5e-4。
  • dataset.batch_size:每张 GPU 的批次大小。默认值为 64。Re-ID 受益于较大的批次,以便更好地进行三元组/对比采样。
  • dataset.num_instances:一个批次中每个身份的实例数量。控制度量学习的采样策略。

多 GPU/多节点

**启动方式:**由 Lightning 管理(单个 python 进程,Lightning 会生成工作进程)。

Spec Key描述默认值
train.num_gpusGPU 数量1
train.gpu_idsGPU 设备索引[0]
  • 多 GPU 策略:ddp_find_unused_parameters_true
  • 始终启用 sync_batchnorm
  • 精度强制设为 FP16(16-mixed)
  • 没有显式的 num_nodes 配置——面向单节点

硬件

至少需要 1 张 GPU,建议使用 2 张 GPU。每张 GPU 需有 16GB 以上显存。Re-ID 模型相对轻量,但度量学习受益于较大的批次大小。

错误模式

num_classes 不匹配:确保 dataset.num_classes 等于训练集中唯一身份文件夹的数量。

Query/gallery 不匹配:Query 和 test(gallery)数据集必须使用相同的身份命名空间。

Spec 参数/父模型推断

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

TAO Core re_identification.config.json 中的推断映射:

操作Spec 字段推断函数含义
evaluateencryption_keykey加密密钥
evaluateevaluate.checkpointparent_model从父作业结果文件夹推断出的模型文件
evaluateevaluate.output_cmc_curve_plotcreate_evaluate_cmc_plot_reidReID CMC 绘图路径
evaluateevaluate.output_sampled_matches_plotcreate_evaluate_matches_plot_reidReID 采样匹配绘图路径
evaluateresults_diroutput_dir当前作业结果目录
exportencryption_keykey加密密钥
exportexport.checkpointparent_model从父作业结果文件夹推断出的模型文件
exportexport.onnx_filecreate_onnx_file输出 ONNX 路径
exportresults_diroutput_dir当前作业结果目录
inferenceencryption_keykey加密密钥
inferenceinference.checkpointparent_model从父作业结果文件夹推断出的模型文件
inferenceinference.output_filecreate_inference_result_file_reidReID 推断 JSON 路径
inferenceresults_diroutput_dir当前作业结果目录
trainencryption_keykey加密密钥
trainmodel.pretrained_model_pathptm_if_no_resume_model不存在恢复检查点时使用的 PTM
trainresults_diroutput_dir当前作业结果目录
traintrain.resume_training_checkpoint_pathresume_model从当前作业结果文件夹推断出的模型文件

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

评论

支持 Markdown