关于 tao-train-reid
用于行人重识别的 TAO 技能,通过学习判别性嵌入,在不同摄像头视角下匹配同一人。支持 AutoML 训练流程及基于检索的评估。
主要功能
- 01跨摄像头行人匹配
- 02判别性嵌入学习
- 03支持 AutoML 的训练流程
- 04CMC 与 mAP 评估
适用场景
- 01跨摄像头查找同一行人
- 02用查询集和图库评估身份匹配
安装
复制下面的命令,在终端中运行。
npx skills add https:/ / github.com/ promptingcompany/ nv-skills --skill tao-train-reid安装后在 Claude Code 中输入/tao-train-reid确认已加载。
包含的文件
除 SKILL.md 外,安装时一起下载的参考资料和脚本。
BENCHMARK.mdskill-card.mdskill.oms.sigevals/evals.jsonreferences/skill_info.yamlspec_template_evaluate.yamlspec_template_export.yamlspec_template_inference.yamlspec_template_train.yamlschemas/evaluate.schema.jsonexport.schema.jsoninference.schema.jsonmanifest.jsontrain.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
各操作的数据集要求
| 操作 | 规格键 | 来源 | 文件 | 列表? |
|---|---|---|---|---|
| evaluate | evaluate.test_dataset | train_datasets | sample_test.tar.gz | 否 |
| evaluate | evaluate.query_dataset | train_datasets | sample_query.tar.gz | 否 |
| inference | inference.test_dataset | train_datasets | sample_test.tar.gz | 否 |
| inference | inference.query_dataset | train_datasets | sample_query.tar.gz | 否 |
| train | dataset.train_dataset_dir | train_datasets | sample_train.tar.gz | 否 |
| train | dataset.test_dataset_dir | train_datasets | sample_test.tar.gz | 否 |
| train | dataset.query_dataset_dir | train_datasets | sample_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_gpus | GPU 数量 | 1 |
train.gpu_ids | GPU 设备索引 | [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 字段 | 推断函数 | 含义 |
|---|---|---|---|
| evaluate | encryption_key | key | 加密密钥 |
| evaluate | evaluate.checkpoint | parent_model | 从父作业结果文件夹推断出的模型文件 |
| evaluate | evaluate.output_cmc_curve_plot | create_evaluate_cmc_plot_reid | ReID CMC 绘图路径 |
| evaluate | evaluate.output_sampled_matches_plot | create_evaluate_matches_plot_reid | ReID 采样匹配绘图路径 |
| 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_reid | ReID 推断 JSON 路径 |
| 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,也不要修改生成的运行器脚本来猜测检查点路径。