Tao Train Reid

NVIDIAOfficialCreated by NVIDIA Corporation

Skill
Tool integration
Data and AI
Apache-2.0
Medium risk

From the author

Person re-identification (ReID). Learns discriminative embeddings to match the same person across different camera views, based on metric learning. Use when training, evaluating, exporting, or running inference for a TAO person re-identification model. Trigger phrases include "train ReID", "person re-identification", "cross-camera person matching", "ReID embeddings", "person re-id".

— the description by NVIDIA Corporation

AutoGineer summary

Written by AutoGineer from the skill’s own files, not by the author

This skill provides a structured interface for managing the lifecycle of person re-identification models within the TAO ecosystem. It utilizes metric learning to generate discriminative embeddings, enabling the identification of the same individual across different camera views. The skill supports training, evaluation, inference, and model export, with built-in support for AutoML to optimize performance metrics like cmc_rank_1.

What it does

Trains person re-identification models using metric learning.
Evaluates model performance using mAP and CMC rank metrics.
Exports trained models to ONNX format for deployment.
Performs inference on test and query datasets.
Supports AutoML for automated hyperparameter optimization.
Manages model checkpoints and configuration overrides.

When to use it

When you need to match individuals across multiple camera feeds.
When developing computer vision pipelines for person tracking.
When you need to optimize re-identification model accuracy via AutoML.
When preparing a model for deployment in a production environment.

Advantages

Enables cross-camera person matching through discriminative embeddings.
Automates hyperparameter tuning via integrated AutoML support.
Provides standardized workflows for training, evaluation, and export.
Supports large-batch training for improved metric learning performance.

Usage examples

“Train a new person re-identification model using the provided training specification.”: The agent initiates the training process using the specified configuration, optimizing for cmc_rank_1.
“Evaluate my re-identification model performance.”: The agent runs the evaluation action against the test and query datasets, reporting mAP and CMC ranks.
“Export my trained re-identification model to ONNX.”: The agent processes the checkpoint and generates an ONNX file for deployment.
“Run inference on my test dataset.”: The agent executes the inference action and produces a JSON file containing the results.

Install

Download the skill as a ZIP:
https://autogineer.com/agent-skills/nvidia/tao-train-reid/download

Unzip it into the cross-client skills folder of your project:

mkdir -p .agents/skills && unzip tao-train-reid.zip -d .agents/skills

Documentation by NVIDIA Corporation

Re-Identification

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).

Person re-identification. Learns discriminative embeddings to match the same person across different camera views. Metric learning based.

Set model.pretrained_model_path for pretrained weights.

Quick Start (docker run)

Docker-native launch — no TAO SDK and no Python on the host. Use the local Docker/platform skill instead when it gives a stricter environment-specific command (non-root UID mapping, cache redirects, remote daemons).

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"
)

Train:

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

Evaluate:

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

Inference:

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

Export:

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

Every action takes its spec with -e; results_dir is set in the spec or overridden on the command line. Mount any pretrained-weights directory the spec references, and keep every in-container path consistent across actions.

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.

Supported Actions

The packaged Re-Identification PyT CLI supports train, evaluate, inference, export, and default_specs. This model skill exposes the runnable user actions train, evaluate, inference, and export; resume/retrain is performed through train with train.resume_training_checkpoint_path.

Do not advertise or synthesize dataset_convert, deploy, prune, quantize, gen_trt_engine, or standalone retrain for this model unless the packaged model skill and real CLI add those actions.

Training Requirements

Dataset type: re_identification
Formats: default
AutoML training metric: cmc_rank_1 (maximize). Train status files emit this retrieval KPI for trial ranking.
Standalone evaluation metric: mAP. The evaluate action writes mAP to its status KPI and prints the CMC ranks in its console table; do not expect standalone evaluate status to contain cmc_rank_1.

Per-Action Dataset Requirements

Action	Spec Key	Source	Files	List?
evaluate	evaluate.test_dataset	train_datasets	sample_test.tar.gz	No
evaluate	evaluate.query_dataset	train_datasets	sample_query.tar.gz	No
inference	inference.test_dataset	train_datasets	sample_test.tar.gz	No
inference	inference.query_dataset	train_datasets	sample_query.tar.gz	No
train	dataset.train_dataset_dir	train_datasets	sample_train.tar.gz	No
train	dataset.test_dataset_dir	train_datasets	sample_test.tar.gz	No
train	dataset.query_dataset_dir	train_datasets	sample_query.tar.gz	No

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"
CHECKPOINT = "/results/{train_job_id}/results_dir/model_epoch_000_step_00099.pth"

train (mandatory data sources):

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

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

{
    "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 (mandatory checkpoint and output):

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

inference (mandatory data sources and checkpoint):

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

For export and inference, provide explicit file paths for export.onnx_file and inference.output_file. For evaluate, provide explicit file paths for evaluate.output_cmc_curve_plot and evaluate.output_sampled_matches_plot. Keep these as spec values or spec_params mappings; do not declare them as file outputs in skill_info.yaml for local Docker until the runner distinguishes files from folders during output pre-creation.

Eval Dataset

Required. Evaluation requires test and query datasets for retrieval-based metrics (CMC, mAP).

Important Parameters

dataset.num_classes: Number of identities. Default 751. Must match the number of unique identities in training data.
model.backbone: Default resnet_50.
optim.base_lr: Base learning rate. Default 3.5e-4.
dataset.batch_size: Per-GPU batch size. Default 64. Re-ID benefits from large batches for better triplet/contrastive sampling.
dataset.num_instances: Number of instances per identity in a batch. Controls sampling strategy for metric learning.

Multi-GPU / Multi-Node

Launch method: Lightning-managed (single python process, Lightning spawns workers).

Spec Key	Description	Default
`train.num_gpus`	Number of GPUs	1
`train.gpu_ids`	GPU device indices	[0]

Multi-GPU strategy: ddp_find_unused_parameters_true
sync_batchnorm is always enabled
Precision forced to FP16 (16-mixed)
No explicit num_nodes config — single-node oriented

Hardware

Minimum 1 GPU(s), recommended 2 GPU(s). 16GB+ VRAM per GPU. Re-ID models are relatively lightweight but benefit from large batch sizes for metric learning.

Error Patterns

num_classes mismatch: Ensure dataset.num_classes equals the number of unique identity folders in the training set.

Invalid triplet batch shape: dataset.batch_size must be compatible with dataset.num_instances so each mini-batch can be reshaped for hard-example mining. For local AutoML smoke runs, keep dataset.batch_size fixed to a known valid multiple such as 16 with dataset.num_instances: 4, and tune train.optim.base_lr instead of unconstrained batch size.

Query/gallery mismatch: Query and test (gallery) datasets must share the same identity namespace.

PyTorch 2.6 checkpoint load failure on checkpoint consumers: Current Re-ID checkpoints include OmegaConf containers. For checkpoints produced by the same trusted TAO train/AutoML workflow, set TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 in downstream resume, evaluate, inference, and export job env vars so Lightning/PyTorch can load the full checkpoint. Do not use this env var for untrusted checkpoints.

AutoML metric extraction: Re-ID train status files report retrieval KPIs such as cmc_rank_1, cmc_rank_5, cmc_rank_10, and mAP, plus train loss. Default AutoML train launches must optimize cmc_rank_1 with direction: maximize; do not use val_loss as the metric for this model.

Checkpoint handoff: Use the checkpoint resolver on the best AutoML child job's results_dir/train/ folder and select the action-appropriate model_epoch_*.pth checkpoint. Re-ID also writes reid_model_latest.pth, but that is a latest symlink and should only be used when a caller explicitly requests latest. Preserve the same dataset identity count and query/gallery archives for downstream actions.

Default spec generation: The packaged default_specs CLI action does not consume the normal -e <spec.yaml> experiment file for results_dir. Invoke it with a Hydra override such as re_identification default_specs results_dir=/workspace/run/results/default_specs. Passing only -e leaves cfg.results_dir unset and fails with MissingMandatoryValue: results_dir.

Spec Param / Parent Model Inference

Model-specific inference mappings belong in this MD file, not in config.json. Generated runners should read this section and apply the mappings with SDK helpers before create_job(). This mirrors the old microservices infer_params.py flow.

Inference mappings from TAO Core re_identification.config.json:

Action	Spec Field	Inference Function	Meaning
evaluate	`encryption_key`	`key`	encryption key
evaluate	`evaluate.checkpoint`	`parent_model`	model file inferred from the parent job results folder
evaluate	`evaluate.output_cmc_curve_plot`	`create_evaluate_cmc_plot_reid`	ReID CMC plot path
evaluate	`evaluate.output_sampled_matches_plot`	`create_evaluate_matches_plot_reid`	ReID sampled matches plot path
evaluate	`results_dir`	`output_dir`	current job results directory
export	`encryption_key`	`key`	encryption key
export	`export.checkpoint`	`parent_model`	model file inferred from the parent job results folder
export	`export.onnx_file`	`create_onnx_file`	output ONNX path
export	`results_dir`	`output_dir`	current job results directory
inference	`encryption_key`	`key`	encryption key
inference	`inference.checkpoint`	`parent_model`	model file inferred from the parent job results folder
inference	`inference.output_file`	`create_inference_result_file_reid`	ReID inference JSON path
inference	`results_dir`	`output_dir`	current job results directory
train	`encryption_key`	`key`	encryption key
train	`model.pretrained_model_path`	`ptm_if_no_resume_model`	PTM when no resume checkpoint exists
train	`results_dir`	`output_dir`	current job results directory
train	`train.resume_training_checkpoint_path`	`resume_model`	model file inferred from the current job results folder

For parent_model or parent_model_folder, pass the upstream train/export/AutoML child job id as parent_job_id. The SDK lists the parent result folder, filters checkpoint artifacts, and returns the selected model file or folder. Do not add these mappings back to config.json and do not patch generated runner scripts to guess checkpoint paths.

Files

ATTRIBUTION.mdAdded by AutoGineer912 B
BENCHMARK.md6.9 KiB
LICENSELicence file added by AutoGineer11.2 KiB
SKILL.md14.1 KiB
config/skillspector-baseline.yaml1.0 KiB
evals/evals.json752 B
references/skill_info.yaml4.4 KiB
references/spec_template_evaluate.yaml2.4 KiB
references/spec_template_export.yaml2.2 KiB
references/spec_template_inference.yaml2.3 KiB
references/spec_template_train.yaml2.2 KiB
schemas/evaluate.schema.json37.4 KiB
schemas/export.schema.json34.9 KiB
schemas/inference.schema.json37.1 KiB
schemas/manifest.json6.5 KiB
schemas/train.schema.json33.7 KiB
skill-card.md4.4 KiB
skill.oms.sig7.1 KiB

schemas/export.schema.json

{
  "automl_default_parameters": [
    "train.optim.lr_steps"
  ],
  "automl_disabled_parameters": [
    "train.cudnn",
    "evaluate",
    "inference",
    "train",
    "re_ranking",
    "train.gpu_ids",
    "export",
    "model",
    "wandb",
    "wandb.tags",
    "dataset",
    "inference.gpu_ids",
    "evaluate.gpu_ids",
    "train.optim",
    "dataset.pixel_std",
    "dataset.pixel_mean",
    "model.stride_size"
  ],
  "default": {
    "dataset": {
      "batch_size": 64,
      "num_classes": 751,
      "num_instances": 4,
      "num_workers": 8,
      "padding": 10,
      "pixel_mean": [
        0.485,
        0.456,
        0.406
      ],
      "pixel_std": [
        0.226,
        0.226,
        0.226
      ],
      "prob": 0.5,
      "query_dataset_dir": "",
      "re_prob": 0.5,
      "sampler": "softmax_triplet",
      "test_dataset_dir": "",
      "train_dataset_dir": "",
      "val_batch_size": 128
    },
    "encryption_key": "",
    "export": {
      "checkpoint": "",
      "gpu_id": 0,
      "onnx_file": "",
      "results_dir": ""
    },
    "model": {
      "att_drop_rate": 0.0,
      "backbone": "resnet_50",
      "cos_layer": false,
      "devide_length": 4,
      "drop_out": 0.0,
      "drop_path": 0.1,
      "dropout_rate": 0.0,
      "feat_dim": 256,
      "gem_pooling": false,
      "id_loss_type": "softmax",
      "id_loss_weight": 1.0,
      "input_channels": 3,
      "input_height": 256,
      "input_width": 128,
      "jpm": false,
      "label_smooth": true,
      "last_stride": 1,
      "metric_loss_type": "triplet",
      "neck": "bnneck",
      "neck_feat": "after",
      "no_margin": false,
      "pretrain_choice": "imagenet",
      "pretrain_hw_ratio": 2.0,
      "pretrained_model_path": "",
      "re_arrange": true,
      "reduce_feat_dim": false,
      "semantic_weight": 1.0,
      "shift_num": 5,
      "shuffle_group": 2,
      "sie_camera": false,
      "sie_coe": 3.0,
      "sie_view": false,
      "stem_conv": false,
      "stride_size": [
        16,
        16
      ],
      "triplet_loss_weight": 1.0,
      "with_center_loss": false,
      "with_flip_feature": false
    },
    "model_name": "",
    "re_ranking": {
      "k1": 20,
      "k2": 6,
      "lambda_value": 0.3,
      "max_rank": 10,
      "num_query": 10,
      "re_ranking": false
    },
    "results_dir": "",
    "train": {
      "checkpoint_interval": 1,
      "checkpoint_interval_unit": "epoch",
      "cudnn": {
        "benchmark": false,
        "deterministic": true
      },
      "gpu_ids": [
        0
      ],
      "grad_clip": 0.0,
      "num_epochs": 10,
      "num_gpus": 1,
      "num_nodes": 1,
      "optim": {
        "base_lr": 0.00035,
        "bias_lr_factor": 1.0,
        "center_loss_weight": 0.0005,
        "center_lr": 0.5,
        "cosine_margin": 0.5,
        "cosine_scale": 30.0,
        "gamma": 0.1,
        "large_fc_lr": false,
        "lr_monitor": "val_loss",
        "lr_steps": [
          40,
          70
        ],
        "momentum": 0.9,
        "name": "Adam",
        "triplet_loss_margin": 0.3,
        "trp_l2": false,
        "warmup_epochs": 20,
        "warmup_factor": 0.01,
        "warmup_iters": 10,
        "warmup_method": "linear",
        "weight_decay": 0.0005,
        "weight_decay_bias": 0.0005
      },
      "results_dir": "",
      "resume_training_checkpoint_path": "",
      "seed": 1234,
      "validation_interval": 1
    },
    "wandb": {
      "enable": true,
      "entity": "",
      "group": "",
      "name": "TAO Toolkit Training",
      "project": "TAO Toolkit",
      "reinit": false,
      "run_id": "",
      "save_code": false,
      "sync_tensorboard": false,
      "tags": [
        "tao-toolkit"
      ]
    }
  },
  "popular": {
    "dataset": {
      "batch_size": 64,
      "num_classes": 751,
      "num_instances": 4,
      "num_workers": 8,
      "val_batch_size": 128
    },
    "evaluate": {
      "gpu_ids": [
        0
      ],
      "num_gpus": 1,
      "num_nodes": 1
    },
    "inference": {
      "gpu_ids": [
        0
      ],
      "num_gpus": 1,
      "num_nodes": 1
    },
    "model": {
      "backbone": "resnet_50"
    },
    "re_ranking": {
      "k1": 20,
      "k2": 6,
      "max_rank": 10,
      "num_query": 10
    },
    "train": {
      "checkpoint_interval": 1,
      "gpu_ids": [
        0
      ],
      "num_epochs": 10,
      "num_gpus": 1,
      "num_nodes": 1,
      "validation_interval": 1
    }
  },
  "properties": {
    "automl_disabled_parameters": [
      "wandb",
      "model",
      "dataset",
      "train",
      "evaluate",
      "inference",
      "export",
      "re_ranking"
    ],
    "dataset": {
      "automl_disabled_parameters": [
        "dataset.pixel_mean",
        "dataset.pixel_std"
      ],
      "automl_enabled": false,
      "default": {
        "batch_size": 64,
        "num_classes": 751,
        "num_instances": 4,
        "num_workers": 8,
        "padding": 10,
        "pixel_mean": [
          0.485,
          0.456,
          0.406
        ],
        "pixel_std": [
          0.226,
          0.226,
          0.226
        ],
        "prob": 0.5,
        "query_dataset_dir": "",
        "re_prob": 0.5,
        "sampler": "softmax_triplet",
        "test_dataset_dir": "",
        "train_dataset_dir": "",
        "val_batch_size": 128
      },
      "description": "Configurable parameters to construct the dataset for a Re-Identification experiment.",
      "popular": [
        "num_instances",
        "val_batch_size",
        "batch_size",
        "num_workers",
        "num_classes"
      ],
      "properties": {
        "batch_size": {
          "default": 64,
          "description": "Batch size.",
          "maximum": Infinity,
          "minimum": 1,
          "popular": true,
          "title": "Batch Size",
          "type": "int"
        },
        "num_classes": {
          "default": 751,
          "description": "Number of classes.",
          "maximum": Infinity,
          "minimum": 1,
          "popular": true,
          "title": "Number of Classes",
          "type": "int"
        },
        "num_instances": {
          "default": 4,
          "description": "Number of instances per class in a batch.",
          "maximum": Infinity,
          "minimum": 4,
          "popular": true,
          "title": "Number of Instances",
          "type": "int"
        },
        "num_workers": {
          "default": 8,
          "description": "Number of workers.",
          "maximum": Infinity,
          "minimum": 1,
          "popular": true,
          "title": "Workers",
          "type": "int"
        },
        "padding": {
          "default": 10,
          "description": "Padding size.",
          "maximum": 10,
          "minimum": 0,
          "title": "Padding",
          "type": "int"
        },
        "pixel_mean": {
          "automl_enabled": false,
          "default": [
            0.485,
            0.456,
            0.406
          ],
          "description": "Mean values for normalization.",
          "title": "Pixel Mean",
          "type": "list"
        },
        "pixel_std": {
          "automl_enabled": false,
          "default": [
            0.226,
            0.226,
            0.226
          ],
          "description": "Standard deviation values for normalization.",
          "title": "Pixel Standard Deviation",
          "type": "list"
        },
        "prob": {
          "default": 0.5,
          "description": "Probability for certain augmentations.",
          "maximum": 1.0,
          "minimum": 0.0,
          "title": "Probability",
          "type": "float"
        },
        "query_dataset_dir": {
          "default": "",
          "description": "Directory for the query dataset.",
          "title": "Query Dataset Directory",
          "type": "string"
        },
        "re_prob": {
          "default": 0.5,
          "description": "Probability for re-augmentation.",
          "maximum": 1.0,
          "minimum": 0.0,
          "title": "Re-augmentation Probability",
          "type": "float"
        },
        "sampler": {
          "default": "softmax_triplet",
          "description": "Type of sampler used for selecting instances.",
          "title": "Sampler Type",
          "type": "string"
        },
        "test_dataset_dir": {
          "default": "",
          "description": "Directory for the testing dataset.",
          "title": "Testing Dataset Directory",
          "type": "string"
        },
        "train_dataset_dir": {
          "default": "",
          "description": "Directory for the training dataset.",
          "title": "Training Dataset Directory",
          "type": "string"
        },
        "val_batch_size": {
          "default": 128,
          "description": "Validation Batch size.",
          "maximum": Infinity,
          "minimum": 1,
          "popular": true,
          "title": "Validation Batch Size",
          "type": "int"
        }
      },
      "type": "collection"
    },
    "encryption_key": {
      "default": "",
      "description": "Key for encrypting model checkpoints",
      "title": "Encryption key",
      "type": "string"
    },
    "export": {
      "automl_enabled": false,
      "default": {
        "checkpoint": "",
        "gpu_id": 0,
        "onnx_file": "",
        "results_dir": ""
      },
      "description": "Configurable parameters to construct the exporter for a Re-Identification experiment.",
      "properties": {
        "checkpoint": {
          "default": "",
          "description": "Path to the checkpoint file.",
          "title": "Checkpoint File",
          "type": "string"
        },
        "gpu_id": {
          "default": 0,
          "description": "GPU ID for computation.",
          "maximum": Infinity,
          "minimum": 0,
          "title": "GPU ID",
          "type": "int"
        },
        "onnx_file": {
          "default": "",
          "description": "Path to the ONNX model file.",
          "title": "ONNX File",
          "type": "string"
        },
        "results_dir": {
          "default": "",
          "description": "Directory for storing results.",
          "title": "Results Directory",
          "type": "string"
        }
      },
      "type": "collection"
    },
    "model": {
      "automl_disabled_parameters": [
        "model.stride_size"
      ],
      "automl_enabled": false,
      "default": {
        "att_drop_rate": 0.0,
        "backbone": "resnet_50",
        "cos_layer": false,
        "devide_length": 4,
        "drop_out": 0.0,
        "drop_path": 0.1,
        "dropout_rate": 0.0,
        "feat_dim": 256,
        "gem_pooling": false,
        "id_loss_type": "softmax",
        "id_loss_weight": 1.0,
        "input_channels": 3,
        "input_height": 256,
        "input_width": 128,
        "jpm": false,
        "label_smooth": true,
        "last_stride": 1,
        "metric_loss_type": "triplet",
        "neck": "bnneck",
        "neck_feat": "after",
        "no_margin": false,
        "pretrain_choice": "imagenet",
        "pretrain_hw_ratio": 2.0,
        "pretrained_model_path": "",
        "re_arrange": true,
        "reduce_feat_dim": false,
        "semantic_weight": 1.0,
        "shift_num": 5,
        "shuffle_group": 2,
        "sie_camera": false,
        "sie_coe": 3.0,
        "sie_view": false,
        "stem_conv": false,
        "stride_size": [
          16,
          16
        ],
        "triplet_loss_weight": 1.0,
        "with_center_loss": false,
        "with_flip_feature": false
      },
      "description": "Configurable parameters to construct the model for a Re-Identification experiment.",
      "popular": [
        "backbone"
      ],
      "properties": {
        "att_drop_rate": {
          "default": 0.0,
          "description": "Attention dropout rate.",
          "maximum": 1.0,
          "minimum": 0.0,
          "title": "Attention Drop Rate",
          "type": "float"
        },
        "backbone": {
          "default": "resnet_50",
          "description": "Backbone type.",
          "popular": true,
          "title": "Backbone Type",
          "type": "string"
        },
        "cos_layer": {
          "default": false,
          "description": "Whether cosine layer is used for the output.",
          "title": "Cosine Layer",
          "type": "bool"
        },
        "devide_length": {
          "default": 4,
          "description": "Length for division in the re-arrangement process.",
          "maximum": 4,
          "minimum": 4,
          "title": "Divide Length",
          "type": "int"
        },
        "drop_out": {
          "default": 0.0,
          "description": "Dropout probability.",
          "maximum": 1.0,
          "minimum": 0.0,
          "title": "Drop Out",
          "type": "float"
        },
        "drop_path": {
          "default": 0.1,
          "description": "Drop path probability.",
          "maximum": 1.0,
          "minimum": 0.0,
          "title": "Drop Path",
          "type": "float"
        },
        "dropout_rate": {
          "default": 0.0,
          "description": "Dropout rate applied in the model.",
          "maximum": 1.0,
          "minimum": 0.0,
          "title": "Dropout Rate",
          "type": "float"
        },
        "feat_dim": {
          "default": 256,
          "description": "Dimension of the feature vector.",
          "maximum": 768,
          "minimum": 32,
          "title": "Feature Dimension",
          "type": "int"
        },
        "gem_pooling": {
          "default": false,
          "description": "Whether generalized mean pooling is used.",
          "title": "GEM Pooling",
          "type": "bool"
        },
        "id_loss_type": {
          "default": "softmax",
          "description": "Type of ID loss used.",
          "title": "ID Loss Type",
          "type": "string"
        },
        "id_loss_weight": {
          "default": 1.0,
          "description": "Weight of the ID loss.",
          "maximum": 1.0,
          "minimum": 0.0,
          "title": "ID Loss Weight",
          "type": "float"
        },
        "input_channels": {
          "default": 3,
          "description": "Number of input channels.",
          "maximum": 3,
          "minimum": 3,
          "title": "Input Channels",
          "type": "int"
        },
        "input_height": {
          "default": 256,
          "description": "Height of the input image.",
          "maximum": 256,
          "minimum": 256,
          "title": "Input Height",
          "type": "int"
        },
        "input_width": {
          "default": 128,
          "description": "Width of the input image.",
          "maximum": 128,
          "minimum": 128,
          "title": "Input Width",
          "type": "int"
        },
        "jpm": {
          "default": false,
          "description": "Whether Joint Part and Global feature learning module is enabled.",
          "title": "JPM",
          "type": "bool"
        },
        "label_smooth": {
          "default": true,
          "description": "Whether label smoothing is applied.",
          "title": "Label Smooth",
          "type": "bool"
        },
        "last_stride": {
          "default": 1,
          "description": "Stride size of the last layer of the backbone.",
          "maximum": 1,
          "minimum": 1,
          "title": "Last Stride",
          "type": "int"
        },
        "metric_loss_type": {
          "default": "triplet",
          "description": "Type of metric loss used.",
          "title": "Metric Loss Type",
          "type": "string"
        },
        "neck": {
          "default": "bnneck",
          "description": "Type of neck used in the model architecture.",
          "title": "Neck Type",
          "type": "string"
        },
        "neck_feat": {
          "default": "after",
          "description": "Position of the feature extraction in the neck.",
          "title": "Neck Feature Position",
          "type": "string"
        },
        "no_margin": {
          "default": false,
          "description": "Whether margin is used in loss computation.",
          "title": "No Margin",
          "type": "bool"
        },
        "pretrain_choice": {
          "default": "imagenet",
          "description": "Source of pretraining.",
          "title": "Pretrain Choice",
          "type": "string"
        },
        "pretrain_hw_ratio": {
          "default": 2.0,
          "description": "Height-width ratio of the pretraining model.",
          "maximum": 2.0,
          "minimum": 2.0,
          "title": "Pretrain HW Ratio",
          "type": "float"
        },
        "pretrained_model_path": {
          "default": "",
          "description": "Path to the pretrained model file.",
          "title": "Pretrained Model Path",
          "type": "string"
        },
        "re_arrange": {
          "default": true,
          "description": "Whether to re-arrange elements in some pattern.",
          "title": "Re-arrange",
          "type": "bool"
        },
        "reduce_feat_dim": {
          "default": false,
          "description": "Whether feature dimension reduction is applied.",
          "title": "Reduce Feature Dimension",
          "type": "bool"
        },
        "semantic_weight": {
          "default": 1.0,
          "description": "Weight for the semantic component in loss calculation.",
          "maximum": 1.0,
          "minimum": 0.0,
          "title": "Semantic Weight",
          "type": "float"
        },
        "shift_num": {
          "default": 5,
          "description": "Number of positions to shift in shift layer.",
          "maximum": 5,
          "minimum": 5,
          "title": "Shift Number",
          "type": "int"
        },
        "shuffle_group": {
          "default": 2,
          "description": "Number of groups for channel shuffling.",
          "maximum": 2,
          "minimum": 2,
          "title": "Shuffle Group",
          "type": "int"
        },
        "sie_camera": {
          "default": false,
          "description": "Whether camera-based Spatial Information Enhancement is used.",
          "title": "SIE Camera",
          "type": "bool"
        },
        "sie_coe": {
          "default": 3.0,
          "description": "Coefficient for scaling in SIE module.",
          "maximum": Infinity,
          "minimum": 0.0,
          "title": "SIE Coefficient",
          "type": "float"
        },
        "sie_view": {
          "default": false,
          "description": "Whether view-based Spatial Information Enhancement is used.",
          "title": "SIE View",
          "type": "bool"
        },
        "stem_conv": {
          "default": false,
          "description": "Whether a convolutional stem is used at the model input.",
          "title": "Stem Convolution",
          "type": "bool"
        },
        "stride_size": {
          "automl_enabled": false,
          "default": [
            16,
            16
          ],
          "description": "Size of stride in the convolution layers.",
          "title": "Stride Size",
          "type": "list"
        },
        "triplet_loss_weight": {
          "default": 1.0,
          "description": "Weight of the triplet loss.",
          "maximum": 1.0,
          "minimum": 0.0,
          "title": "Triplet Loss Weight",
          "type": "float"
        },
        "with_center_loss": {
          "default": false,
          "description": "Whether center loss is used.",
          "title": "Center Loss",
          "type": "bool"
        },
        "with_flip_feature": {
          "default": false,
          "description": "Whether flip feature is enabled.",
          "title": "Flip Feature",
          "type": "bool"
        }
      },
      "type": "collection"
    },
    "model_name": {
      "default": "",
      "description": "Name of model if invoking task via :code:`model_agnostic`",
      "title": "Model name",
      "type": "string"
    },
    "re_ranking": {
      "automl_enabled": false,
      "default": {
        "k1": 20,
        "k2": 6,
        "lambda_value": 0.3,
        "max_rank": 10,
        "num_query": 10,
        "re_ranking": false
      },
      "description": "Configurable parameters to construct the re-ranking parameters for a Re-Identification experiment.",
      "popular": [
        "max_rank",
        "k1",
        "num_query",
        "k2"
      ],
      "properties": {
        "k1": {
          "default": 20,
          "description": "The number of top-k candidates in the first round of re-ranking.",
          "maximum": 20,
          "minimum": 20,
          "popular": true,
          "title": "K1",
          "type": "int"
        },
        "k2": {
          "default": 6,
          "description": "The number of top-k candidates in the second round of re-ranking.",
          "maximum": 6,
          "minimum": 6,
          "popular": true,
          "title": "K2",
          "type": "int"
        },
        "lambda_value": {
          "default": 0.3,
          "description": "The lambda value for balancing the original and Jaccard distance in re-ranking.",
          "maximum": 0.3,
          "minimum": 0.0,
          "title": "Lambda Value",
          "type": "float"
        },
        "max_rank": {
          "default": 10,
          "description": "The maximum rank considered in re-ranking.",
          "maximum": 10,
          "minimum": 10,
          "popular": true,
          "title": "Max Rank",
          "type": "int"
        },
        "num_query": {
          "default": 10,
          "description": "The number of query images used in re-ranking.",
          "maximum": 10,
          "minimum": 10,
          "popular": true,
          "title": "Number of Queries",
          "type": "int"
        },
        "re_ranking": {
          "default": false,
          "description": "Enable or disable re-ranking.",
          "title": "Re-Ranking",
          "type": "bool"
        }
      },
      "type": "collection"
    },
    "results_dir": {
      "default": "",
      "description": "\n        Path to where all the assets generated from a task are stored.\n        ",
      "title": "Results directory",
      "type": "string"
    },
    "train": {
      "automl_disabled_parameters": [
        "train.gpu_ids",
        "train.cudnn",
        "train.optim"
      ],
      "automl_enabled": false,
      "default": {
        "checkpoint_interval": 1,
        "checkpoint_interval_unit": "epoch",
        "cudnn": {
          "benchmark": false,
          "deterministic": true
        },
        "gpu_ids": [
          0
        ],
        "grad_clip": 0.0,
        "num_epochs": 10,
        "num_gpus": 1,
        "num_nodes": 1,
        "optim": {
          "base_lr": 0.00035,
          "bias_lr_factor": 1.0,
          "center_loss_weight": 0.0005,
          "center_lr": 0.5,
          "cosine_margin": 0.5,
          "cosine_scale": 30.0,
          "gamma": 0.1,
          "large_fc_lr": false,
          "lr_monitor": "val_loss",
          "lr_steps": [
            40,
            70
          ],
          "momentum": 0.9,
          "name": "Adam",
          "triplet_loss_margin": 0.3,
          "trp_l2": false,
          "warmup_epochs": 20,
          "warmup_factor": 0.01,
          "warmup_iters": 10,
          "warmup_method": "linear",
          "weight_decay": 0.0005,
          "weight_decay_bias": 0.0005
        },
        "results_dir": "",
        "resume_training_checkpoint_path": "",
        "seed": 1234,
        "validation_interval": 1
      },
      "description": "Configurable parameters to construct the trainer for a Re-Identification experiment.",
      "popular": [
        "validation_interval",
        "checkpoint_interval",
        "num_gpus",
        "num_nodes",
        "num_epochs",
        "gpu_ids"
      ],
      "properties": {
        "checkpoint_interval": {
          "default": 1,
          "description": "The interval (in epochs) at which a checkpoint will be saved. Helps resume training.",
          "minimum": 1,
          "popular": true,
          "title": "Checkpoint interval",
          "type": "int"
        },
        "checkpoint_interval_unit": {
          "default": "epoch",
          "description": "The unit of the checkpoint interval.",
          "enum": [
            "epoch",
            "step"
          ],
          "title": "Checkpoint interval unit",
          "type": "categorical"
        },
        "cudnn": {
          "automl_enabled": false,
          "default": {
            "benchmark": false,
            "deterministic": true
          },
          "properties": {
            "benchmark": {
              "default": false,
              "type": "bool"
            },
            "deterministic": {
              "default": true,
              "type": "bool"
            }
          },
          "type": "collection"
        },
        "gpu_ids": {
          "automl_enabled": false,
          "default": [
            0
          ],
          "description": "\n        List of GPU IDs to run the training on. The length of this list\n        must be equal to the number of gpus in train.num_gpus.",
          "popular": true,
          "title": "GPU IDs",
          "type": "list"
        },
        "grad_clip": {
          "default": 0.0,
          "description": "Maximum norm of the gradients for clipping.",
          "maximum": Infinity,
          "minimum": 0.0,
          "title": "Gradient Clipping",
          "type": "float"
        },
        "num_epochs": {
          "default": 10,
          "description": "Number of epochs to run the training.",
          "maximum": Infinity,
          "minimum": 1,
          "popular": true,
          "title": "Number of epochs",
          "type": "int"
        },
        "num_gpus": {
          "default": 1,
          "description": "The number of GPUs to run the train job.",
          "minimum": 1,
          "popular": true,
          "title": "Number of GPUs",
          "type": "int"
        },
        "num_nodes": {
          "default": 1,
          "description": "Number of nodes to run the training on. If > 1, then multi-node is enabled.",
          "minimum": 1,
          "popular": true,
          "title": "Number of nodes",
          "type": "int"
        },
        "optim": {
          "automl_default_parameters": [
            "train.optim.lr_steps"
          ],
          "automl_enabled": false,
          "default": {
            "base_lr": 0.00035,
            "bias_lr_factor": 1.0,
            "center_loss_weight": 0.0005,
            "center_lr": 0.5,
            "cosine_margin": 0.5,
            "cosine_scale": 30.0,
            "gamma": 0.1,
            "large_fc_lr": false,
            "lr_monitor": "val_loss",
            "lr_steps": [
              40,
              70
            ],
            "momentum": 0.9,
            "name": "Adam",
            "triplet_loss_margin": 0.3,
            "trp_l2": false,
            "warmup_epochs": 20,
            "warmup_factor": 0.01,
            "warmup_iters": 10,
            "warmup_method": "linear",
            "weight_decay": 0.0005,
            "weight_decay_bias": 0.0005
          },
          "description": "Training optimization config.",
          "properties": {
            "base_lr": {
              "default": 0.00035,
              "description": "Base learning rate.",
              "maximum": 1.0,
              "minimum": 0.0,
              "title": "Base Learning Rate",
              "type": "float"
            },
            "bias_lr_factor": {
              "default": 1.0,
              "description": "Learning rate factor for bias parameters.",
              "maximum": 1.0,
              "minimum": 0.0,
              "title": "Bias LR Factor",
              "type": "float"
            },
            "center_loss_weight": {
              "default": 0.0005,
              "description": "Weight of the center loss in the loss function.",
              "maximum": 1.0,
              "minimum": 0.0,
              "title": "Center Loss Weight",
              "type": "float"
            },
            "center_lr": {
              "default": 0.5,
              "description": "Learning rate for center loss parameters.",
              "maximum": 1.0,
              "minimum": 0.0,
              "title": "Center Learning Rate",
              "type": "float"
            },
            "cosine_margin": {
              "default": 0.5,
              "description": "Margin for cosine similarity in losses.",
              "maximum": 1.0,
              "minimum": 0.0,
              "title": "Cosine Margin",
              "type": "float"
            },
            "cosine_scale": {
              "default": 30.0,
              "description": "Scaling factor for cosine similarity.",
              "maximum": Infinity,
              "minimum": 1.0,
              "title": "Cosine Scale",
              "type": "float"
            },
            "gamma": {
              "default": 0.1,
              "description": "Factor by which the learning rate will decrease.",
              "maximum": 1.0,
              "minimum": 0.0,
              "title": "LR Decay Factor",
              "type": "float"
            },
            "large_fc_lr": {
              "default": false,
              "description": "Use a larger learning rate for the fully connected layer.",
              "title": "Large FC Learning Rate",
              "type": "bool"
            },
            "lr_monitor": {
              "default": "val_loss",
              "description": "Metric to monitor for learning rate adjustments.",
              "title": "LR Monitor Metric",
              "type": "string"
            },
            "lr_steps": {
              "automl_enabled": true,
              "default": [
                40,
                70
              ],
              "description": "Epochs at which the learning rate will decrease.",
              "title": "LR Decay Steps",
              "type": "list_2"
            },
            "momentum": {
              "default": 0.9,
              "description": "Momentum factor for optimization.",
              "maximum": 1.0,
              "minimum": 0.0,
              "title": "Momentum",
              "type": "float"
            },
            "name": {
              "default": "Adam",
              "description": "Name of the optimizer.",
              "title": "Optimizer Name",
              "type": "string"
            },
            "triplet_loss_margin": {
              "default": 0.3,
              "description": "Margin for triplet loss.",
              "maximum": 1.0,
              "minimum": 0.0,
              "title": "Triplet Loss Margin",
              "type": "float"
            },
            "trp_l2": {
              "default": false,
              "description": "Apply L2 normalization in triplet loss calculation.",
              "title": "Triplet L2 Normalization",
              "type": "bool"
            },
            "warmup_epochs": {
              "default": 20,
              "description": "Number of epochs for warm-up.",
              "maximum": Infinity,
              "minimum": 0,
              "title": "Warmup Epochs",
              "type": "int"
            },
            "warmup_factor": {
              "default": 0.01,
              "description": "Initial learning rate as a factor of the base learning rate during warm-up.",
              "maximum": 1.0,
              "minimum": 0.0,
              "title": "Warmup Factor",
              "type": "float"
            },
            "warmup_iters": {
              "default": 10,
              "description": "Number of iterations for warm-up.",
              "maximum": Infinity,
              "minimum": 0,
              "title": "Warmup Iterations",
              "type": "int"
            },
            "warmup_method": {
              "default": "linear",
              "description": "Method used for warm-up (e.g., 'linear', 'exp').",
              "title": "Warmup Method",
              "type": "string"
            },
            "weight_decay": {
              "default": 0.0005,
              "description": "Weight decay for regularization.",
              "maximum": 1.0,
              "minimum": 0.0,
              "title": "Weight Decay",
              "type": "float"
            },
            "weight_decay_bias": {
              "default": 0.0005,
              "description": "Weight decay for bias regularization.",
              "maximum": 1.0,
              "minimum": 0.0,
              "title": "Weight Decay for Bias",
              "type": "float"
            }
          },
          "title": "Optimization config",
          "type": "collection"
        },
        "results_dir": {
          "default": "",
          "description": "\n        Path to where all the assets generated from a task are stored.\n        ",
          "title": "Results directory",
          "type": "string"
        },
        "resume_training_checkpoint_path": {
          "default": "",
          "description": "Path to the checkpoint to resume training from.",
          "title": "Resume checkpoint path",
          "type": "string"
        },
        "seed": {
          "default": 1234,
          "description": "The seed for the initializer in PyTorch. If < 0, disable fixed seed.",
          "maximum": Infinity,
          "minimum": -1,
          "title": "Seed for randomization",
          "type": "int"
        },
        "validation_interval": {
          "default": 1,
          "description": "\n        The interval (in epochs) at which a evaluation\n        will be triggered on the validation dataset.",
          "minimum": 1,
          "popular": true,
          "title": "Validation interval",
          "type": "int"
        }
      },
      "type": "collection"
    },
    "wandb": {
      "automl_disabled_parameters": [
        "wandb.tags"
      ],
      "automl_enabled": false,
      "default": {
        "enable": true,
        "entity": "",
        "group": "",
        "name": "TAO Toolkit Training",
        "project": "TAO Toolkit",
        "reinit": false,
        "run_id": "",
        "save_code": false,
        "sync_tensorboard": false,
        "tags": [
          "tao-toolkit"
        ]
      },
      "properties": {
        "enable": {
          "default": true,
          "type": "bool"
        },
        "entity": {
          "default": "",
          "type": "string"
        },
        "group": {
          "default": "",
          "type": "string"
        },
        "name": {
          "default": "TAO Toolkit Training",
          "type": "string"
        },
        "project": {
          "default": "TAO Toolkit",
          "type": "string"
        },
        "reinit": {
          "default": false,
          "type": "bool"
        },
        "run_id": {
          "default": "",
          "type": "string"
        },
        "save_code": {
          "default": false,
          "type": "bool"
        },
        "sync_tensorboard": {
          "default": false,
          "type": "bool"
        },
        "tags": {
          "automl_enabled": false,
          "default": [
            "tao-toolkit"
          ],
          "type": "list"
        }
      },
      "type": "collection"
    }
  },
  "type": "object",
  "x_tao_schema": {
    "action": "export",
    "core_module": "re_identification",
    "model": "re-identification",
    "network_arch": "re_identification",
    "schema_action": "export",
    "schema_version": 1,
    "source": "tao-core dataclass config"
  }
}

Requirements and compatibility

Works with: Claude Code, Claude apps, OpenAI Codex, GitHub Copilot, Cursor, Gemini CLI, Kiro, OpenCode, Goose, Amp, Roo Code, JetBrains Junie, Factory Droid, Devin (ex-Windsurf), Mistral Vibe, OpenHands, Letta

“Works with” means the tool reads the open Agent Skills format; a skill written for one tool may use features others lack.

Needs: Declares allowed tools, Needs a command-line tool, Needs Python, Needs Node.js

The author says: Requires docker + nvidia-container-toolkit.

Security

What our scanner noticed· Medium risk

Large base64-encoded blobs

A static scan is not a guarantee. Skills are third-party content: review a skill before you let an agent use it.

Attribution and licence

Created by

NVIDIA Corporation

Publisher

NVIDIA (Official)

Source

NVIDIA/skills · skills/tao-train-reid · commit fce3cca

Licence

Apache License 2.0 (redistribution allowed)

LICENSE (added by AutoGineer from LICENSE-APACHE)
ATTRIBUTION.md (written by AutoGineer)

Synced

2026-11-09

AutoGineer is not affiliated with or endorsed by NVIDIA Corporation or NVIDIA.

Download ZIP

The ZIP holds the skill’s files, its licence files and an ATTRIBUTION.md. Download tao-train-reid.zip