tao-train-reid

重新识别

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

行人重新识别。通过学习具有判别性的嵌入表示，匹配不同摄像头视角下的同一行人。基于度量学习。

设置 model.pretrained_model_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" \
  re_identification train -e /specs/train.yaml

评估：

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

推理：

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

导出：

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

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

Dataclass Schema

生成的 TAO Core schema 随包提供，位于 schemas/<action>.schema.json；可用操作列于 schemas/manifest.json。每个生成的 schema 还会根据 schema 顶层的 default 字段生成 references/spec_template_<action>.yaml。AutoML 是否启用由模型层的 references/skill_info.yaml 中的 automl_enabled 声明。某项操作要能运行 AutoML，必须存在且可解析 schemas/<action>.schema.json 和 references/spec_template_<action>.yaml。请使用已打包的所选操作 schema，以获取 automl_default_parameters、automl_disabled_parameters、默认值、最小值/最大值范围、枚举、选项权重、数学条件、依赖项和常用参数。不要假设运行时存在 ~/tao-core；维护者会在打包 skill bank 前重新生成 schema 和模板。

训练操作策略

此模型在模型层启用了 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 或缺少已打包的训练 schema/模板时，才直接进行模型训练；如果缺少 schema，请说明此模型已启用 AutoML，但在生成 schema 前无法运行 AutoML。

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

支持的操作

已打包的 Re-Identification PyT CLI 支持 train、evaluate、inference、export 和 default_specs。此模型技能开放可运行的用户操作 train、evaluate、inference 和 export；通过 train.resume_training_checkpoint_path 使用 train 执行恢复/重新训练。

除非已打包的模型技能和实际 CLI 添加了相应操作，否则不要为此模型宣传或合成 dataset_convert、deploy、prune、quantize、gen_trt_engine 或独立的 retrain。

训练要求

• 数据集类型： re_identification
• 格式： 默认
• AutoML 训练指标： cmc_rank_1（最大化）。训练状态文件会输出此检索 KPI，用于试验排序。
• 独立评估指标： mAP。evaluate 操作会将 mAP 写入状态 KPI，并在控制台表格中打印 CMC 排名；不要期待独立 evaluate 状态中包含 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"
CHECKPOINT = "/results/{train_job_id}/results_dir/model_epoch_000_step_00099.pth"

train（必需数据源）：

{
    "train.num_epochs": 30,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "train.num_gpus": 1,
    "dataset.num_classes": 100,
    "dataset.num_workers": 4,
    "dataset.batch_size": 16,
    "dataset.num_instances": 4,
    "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",
}

resume train（必需检查点）：

{
    "train.num_epochs": 31,
    "train.resume_training_checkpoint_path": CHECKPOINT,
    "dataset.num_classes": 100,
    "dataset.batch_size": 16,
    "dataset.num_instances": 4,
    "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",
    "evaluate.checkpoint": CHECKPOINT,
    "evaluate.output_cmc_curve_plot": "/results/{evaluate_job_id}/results_dir/cmc_curve.png",
    "evaluate.output_sampled_matches_plot": "/results/{evaluate_job_id}/results_dir/sampled_matches.png",
}

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

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

inference（必需数据源和检查点）：

{
    "inference.test_dataset": f"{S3_TRAIN}/sample_test.tar.gz",
    "inference.query_dataset": f"{S3_TRAIN}/sample_query.tar.gz",
    "inference.checkpoint": CHECKPOINT,
    "inference.output_file": "/results/{inference_job_id}/results_dir/reid_inference.json",
}

对于 export 和 inference，请为 export.onnx_file 和 inference.output_file 提供明确的文件路径。对于 evaluate，请为 evaluate.output_cmc_curve_plot 和 evaluate.output_sampled_matches_plot 提供明确的文件路径。将这些路径保留为规格值或 spec_params 映射；在本地 Docker 的 skill_info.yaml 中，请勿将它们声明为文件输出，除非运行器能够在预先创建输出时区分文件和文件夹。

Eval 数据集

必需。评估需要 test 和 query 数据集，以计算基于检索的指标（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 会生成工作进程）。

规格键	描述	默认值
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 等于训练集中唯一身份文件夹的数量。

三元组批次形状无效：dataset.batch_size 必须与 dataset.num_instances 兼容，以便每个小批次都能重塑为适用于困难样本挖掘的形状。对于本地 AutoML 冒烟测试，将 dataset.batch_size 固定为已知有效的倍数，例如 16，并设置 dataset.num_instances: 4；调整 train.optim.base_lr，而不要不受约束地调整批次大小。

查询集／图库集不匹配：查询集和测试集（图库集）必须使用相同的身份命名空间。

PyTorch 2.6 检查点使用方加载失败：当前的 Re-ID 检查点包含 OmegaConf 容器。对于由同一可信 TAO 训练／AutoML 工作流生成的检查点，请在下游恢复、评估、推理和导出作业的环境变量中设置 TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1，以便 Lightning/PyTorch 加载完整检查点。切勿对不可信的检查点使用此环境变量。

AutoML 指标提取：Re-ID 训练状态文件会报告检索 KPI，例如 cmc_rank_1、cmc_rank_5、cmc_rank_10 和 mAP，以及训练损失。默认的 AutoML 训练启动必须以 cmc_rank_1 为优化目标，并设置 direction: maximize；此模型不要使用 val_loss 作为指标。

检查点交接：在最佳 AutoML 子作业的 results_dir/train/ 文件夹上使用检查点解析器，并选择适用于相应操作的 model_epoch_*.pth 检查点。Re-ID 也会写入 reid_model_latest.pth，但它是指向最新版本的符号链接，仅应在调用方明确请求最新版本时使用。下游操作中应保留相同的数据集身份数量以及查询集／图库集存档。

默认配置生成：打包的 default_specs CLI 操作不会使用常规的 -e <spec.yaml> 实验文件来设置 results_dir。 请使用 Hydra 覆盖项调用，例如 re_identification default_specs results_dir=/workspace/run/results/default_specs。 仅传入 -e 会导致 cfg.results_dir 未设置，并报错 MissingMandatoryValue: results_dir。

规格参数 / 父模型推断

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

来自 TAO Core re_identification.config.json 的推断映射：

操作	规格字段	推断函数	含义
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，也不要修改生成的运行器脚本来猜测检查点路径。