Tao Train Metric Learning Recognition

ML Recog

Standalone install? If this session was not initialized by the TAO skill bank plugin, run the tao-setup skill first (host preflight, credentials, cross-skill discovery).

Metric learning recognition for fine-grained visual recognition. Learns embeddings for retrieval-based matching (e.g., retail product recognition). Uses triplet/contrastive losses.

Set model.pretrained_model_path for pretrained backbone.

For TAO Deploy TensorRT actions (gen_trt_engine, TensorRT evaluate, and TensorRT inference), read references/tao-deploy-metric-learning-recognition.md first. Deploy spec templates live in this skill's references/ folder with the spec_template_deploy_*.yaml prefix.

Dataclass Schemas

Generated TAO Core schemas are packaged in schemas/<action>.schema.json, with schemas/manifest.json listing available actions. Each generated schema also emits references/spec_template_<action>.yaml from the schema top-level default field. AutoML enablement is declared at the model layer in references/skill_info.yaml via automl_enabled. Runnable AutoML for an action requires schemas/<action>.schema.json and references/spec_template_<action>.yaml to exist and parse. Use the packaged selected-action schema for automl_default_parameters, automl_disabled_parameters, defaults, min/max bounds, enums, option weights, math conditions, dependencies, and popular parameters. Do not expect ~/tao-core at runtime; maintainers regenerate schemas/templates before packaging the skill bank.

Train Action Policy

This model is AutoML-enabled at the model layer. Before handling any train-stage request, read references/skill_info.yaml and resolve the run override from either an explicit automl_policy value or the user's workflow request. Use automl_policy: on by default and only expose on / off in new launch prompts. Treat phrases like "turn off AutoML", "disable AutoML", "no HPO", or "plain training" as automl_policy: off for this run only. When automl_policy: on, automl_enabled: true, and both schemas/train.schema.json and references/spec_template_train.yaml are packaged, route the train action through tao-skill-bank:tao-run-automl by default with this model's skill_dir. Preserve workflow/application overrides for datasets, specs, output directories, GPU/platform settings, parent checkpoints, and automl_policy. Use direct model training only when automl_policy: off or the packaged train schema/template is missing; in the missing-schema case, report that AutoML is enabled but not runnable for this model until schemas are generated.

Non-train actions such as evaluate, inference, export, and deploy flows stay in this model skill. The per-run automl_policy override does not change model metadata.

Training Requirements

• Dataset type: ml_recog
• Formats: default
• Monitoring metric: val Precision at Rank 1
• AutoML metric contract: Use val Precision at Rank 1 emitted during training and maximize it. Use test Precision at Rank 1 only for standalone selected-checkpoint validation.
• Standalone evaluation metric: test Precision at Rank 1. Use the training val Precision at Rank 1 KPI for AutoML recommendation ranking and the test KPI only to verify the selected checkpoint on the evaluation reference/query split.

Per-Action Dataset Requirements

Typical Spec Overrides

Data source overrides are mandatory for every action — the agent MUST construct data source paths from the Per-Action Dataset Requirements table above and include them in spec_overrides.

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

train (mandatory data sources):

{
    "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 (mandatory data sources):

{
    "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 (mandatory data sources):

{
    "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",
}

Eval Dataset

Required. Evaluation requires reference and query datasets for retrieval metrics.

Important Parameters

• model.backbone: Default resnet_50. Options: resnet_50, resnet_101, fan_small, fan_base, fan_large, fan_tiny, nvdinov2_vit_large_legacy.
• model.feat_dim: Embedding dimension. Default 256. Output feature vector size for similarity matching.
• train.batch_size: Per-GPU batch size. Default 4. val_batch_size also 4. For training and AutoML search, train.batch_size must be divisible by dataset.num_instance.
• dataset.num_instance: Instances per identity in a batch (P/K sampling). Default 4. Controls how many images of the same class appear together. If using a custom AutoML range for train.batch_size, use explicit options that are multiples of this value.
• train.optim.trunk.base_lr: Learning rate for the trunk (backbone). Default 3.5e-4 (Adam).
• train.optim.embedder.base_lr: Learning rate for the embedding head. Default 3.5e-4.
• train.optim.triplet_loss_margin: Margin for triplet loss. Default 0.3. smooth_loss=True by default.
• train.optim.miner_function_margin: Hard mining margin. Default 0.1. Controls pair mining difficulty.
• train.optim.steps: LR decay steps. Default [40, 70] with gamma=0.1.
• dataset.train_dataset: Path to training images organized in class folders.
• dataset.val_dataset: Dict with 'reference' and 'query' keys pointing to ImageNet-format directories for retrieval evaluation.