tao-train-metric-learning-recognition

ML Recog

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

用于细粒度视觉识别的度量学习识别。学习用于基于检索进行匹配的嵌入（例如，零售商品识别）。使用三元组／对比损失。

为预训练骨干网络设置 model.pretrained_model_path。

对于 TAO Deploy TensorRT 操作（gen_trt_engine、TensorRT evaluate 和 TensorRT inference），请先阅读 references/tao-deploy-metric-learning-recognition.md。部署规范模板位于此 skill 的 references/ 文件夹中，文件名前缀为 spec_template_deploy_*.yaml。

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。保留工作流／应用对数据集、规范、输出目录、GPU／平台设置、父检查点和 automl_policy 的覆盖设置。仅当 automl_policy: off 或缺少已打包的训练架构／模板时，才直接进行模型训练；若缺少架构／模板，请说明此模型已启用 AutoML，但在生成架构前无法运行 AutoML。

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

训练要求

• 数据集类型： ml_recog
• 格式： 默认
• 监控指标： val Precision at Rank 1
• AutoML 指标约定： 使用训练期间生成的 val Precision at Rank 1 并将其最大化。仅在单独验证所选检查点时使用 test Precision at Rank 1。
• 单独评估指标： test Precision at Rank 1。使用训练期间的 val Precision at Rank 1 KPI 对 AutoML 推荐结果进行排序，仅使用 test KPI 验证评估参考／查询划分上的所选检查点。

各操作的数据集要求

操作	Spec 键	来源	文件	列表？
evaluate	dataset.val_dataset	train_datasets	reference: metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/reference.tar.gz, query: metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/test.tar.gz	否
inference	dataset.val_dataset	train_datasets	reference: metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/reference.tar.gz, query:	否
inference	inference.input_path	train_datasets	metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/test.tar.gz	否
train	dataset.train_dataset	train_datasets	metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes/train.tar.gz	否
train	dataset.val_dataset	train_datasets	reference: metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes/reference.tar.gz, query: metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes/val.tar.gz	否

常见 Spec 覆盖项

每项操作都必须覆盖数据源——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,
    "dataset.train_dataset": f"{S3_TRAIN}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes/train.tar.gz",
    "dataset.val_dataset": {"reference": f"{S3_TRAIN}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes/reference.tar.gz", "query": f"{S3_TRAIN}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes/val.tar.gz"},
}

evaluate（必需数据源）：

{
    "evaluate.checkpoint": "<selected train/AutoML checkpoint>",
    "dataset.val_dataset": {"reference": f"{S3_TRAIN}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/reference.tar.gz", "query": f"{S3_TRAIN}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/test.tar.gz"},
}

inference（必需数据源）：

{
    "inference.checkpoint": "<selected train/AutoML checkpoint>",
    "dataset.val_dataset": {"reference": f"{S3_TRAIN}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/reference.tar.gz"},
    "inference.input_path": f"{S3_TRAIN}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/test.tar.gz",
}

评估数据集

必需。评估需要 reference 和 query 数据集来计算检索指标。

重要参数

• model.backbone：默认为 resnet_50。选项：resnet_50、resnet_101、fan_small、fan_base、fan_large、fan_tiny、nvdinov2_vit_large_legacy。
• model.feat_dim：嵌入维度。默认为 256。用于相似度匹配的输出特征向量大小。
• train.batch_size：每个 GPU 的批次大小。默认为 4。val_batch_size 也为 4。对于训练和 AutoML 搜索，train.batch_size 必须能被 dataset.num_instance 整除。
• dataset.num_instance：批次中每个身份对应的实例数（P/K 采样）。默认为 4。控制同一类别的多少张图像会出现在一起。如果为 train.batch_size 使用自定义 AutoML 范围，请使用该值的整数倍作为显式选项。
• train.optim.trunk.base_lr：主干网络（骨干网络）的学习率。默认为 3.5e-4（Adam）。
• train.optim.embedder.base_lr：嵌入头的学习率。默认为 3.5e-4。
• train.optim.triplet_loss_margin：三元组损失的边界。默认为 0.3。默认启用 smooth_loss=True。
• train.optim.miner_function_margin：困难样本挖掘边界。默认为 0.1。用于控制样本对挖掘难度。
• train.optim.steps：学习率衰减步数。默认为 [40, 70]，gamma=0.1。
• dataset.train_dataset：按类别文件夹组织的训练图像路径。
• dataset.val_dataset：包含 'reference' 和 'query' 键的字典，指向用于检索评估的 ImageNet 格式目录。

多 GPU / 多节点

**启动方式：**由 Lightning 管理（单个 python 进程，Lightning 会启动工作进程）。

Spec 键	描述	默认值
train.num_gpus	GPU 数量	1
train.gpu_ids	GPU 设备索引	[0]

• 策略：auto（Lightning 自动选择最佳策略）
• 没有显式的 num_nodes 或 distributed_strategy 配置——面向单节点

硬件

至少需要 1 个 GPU，建议使用 2 个 GPU。每个 GPU 需配备 16GB 以上显存。度量学习可以通过较大的批次获得更好的三元组采样效果，但除此之外内存需求适中。

错误模式

参考数据集/查询数据集不匹配：确保参考数据集和查询数据集在评估时使用兼容的类别命名空间。

检查点操作中 PyTorch 2.6 加载检查点失败：当前 TAO ML-Recog 检查点可能包含 OmegaConf 对象。对于由同一可信 TAO 训练/AutoML 工作流生成的检查点，请在下游评估、推理、导出或恢复/重新训练作业的环境变量中设置 TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1，以便 Lightning 加载完整检查点。不要对不可信的检查点使用此环境变量。

Spec 参数 / 父模型推理

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

TAO Core ml_recog.config.json 中的推理映射：

操作	Spec 字段	推理函数	含义
评估	evaluate.checkpoint	parent_model	从父作业结果文件夹推断模型文件
评估	results_dir	output_dir	当前作业结果目录
导出	export.checkpoint	parent_model	从父作业结果文件夹推断模型文件
导出	export.onnx_file	create_onnx_file	输出 ONNX 路径
导出	results_dir	output_dir	当前作业结果目录
推理	inference.checkpoint	parent_model	从父作业结果文件夹推断模型文件
推理	results_dir	output_dir	当前作业结果目录
训练	model.pretrained_model_path	ptm_if_no_resume_model	不存在恢复检查点时使用的 PTM
训练	results_dir	output_dir	当前作业结果目录
训练	train.resume_training_checkpoint_path	resume_model	从当前作业结果文件夹推断模型文件

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

部署

• tao-deploy-metric-learning-recognition