tao-train-reid について
tao-train-reid は、カメラの異なる視点間で同一人物を照合するための TAO スキルです。メトリック学習による埋め込みの訓練に対応し、AutoML を使った訓練のほか、評価、推論、エクスポートを行えます。
主な機能
- 01カメラ間の人物照合
- 02メトリック学習による埋め込み訓練
- 03訓練、評価、推論、エクスポート
- 04AutoML 訓練への対応
- 05Docker コンテナでの実行
活用シーン
- 01カメラ間の再識別モデルの訓練
- 02検索 mAP・CMC 指標の評価
- 03クエリとテストデータでの推論
- 04モデルの ONNX ファイルへのエクスポート
インストール
コマンドをコピーしてターミナルで実行します。
npx skills add https:/ / github.com/ nvidia/ skills --skill tao-train-reidインストール後、Claude Code で/tao-train-reidを入力して読み込まれたか確認します。
同梱ファイル
SKILL.md と一緒にインストールされる参考資料やスクリプト。
BENCHMARK.mdskill-card.mdskill.oms.sigconfig/skillspector-baseline.yamlevals/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詳細
再識別
単体でインストールする場合 このセッションが 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 スキーマ
生成された TAO Core スキーマは schemas/<action>.schema.json にパッケージ化され、利用可能なアクションは schemas/manifest.json に記載されています。各生成スキーマでは、スキーマ最上位の default フィールドから references/spec_template_<action>.yaml も生成されます。AutoML の有効化は、automl_enabled を使って references/skill_info.yaml のモデルレイヤーで宣言します。あるアクションで 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 の両方がパッケージ化されている場合、デフォルトでこのモデルの skill_dir を指定して tao-skill-bank:tao-run-automl 経由で train アクションを実行してください。データセット、spec、出力ディレクトリ、GPU/プラットフォームの設定、親チェックポイント、automl_policy に対するワークフロー/アプリケーションの上書きは維持してください。直接モデルをトレーニングするのは、automl_policy: off の場合、またはパッケージ化された train スキーマ/テンプレートがない場合に限ります。スキーマがない場合は、AutoML は有効ですが、スキーマが生成されるまではこのモデルで実行できないことを報告してください。
evaluate、inference、export などの train 以外のアクションやデプロイフローは、このモデルスキル内で実行します。実行ごとの 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
- 形式: default
- AutoML のトレーニング指標:
cmc_rank_1(最大化)。トレーニングのステータスファイルには、試行のランキング用にこの検索 KPI が出力されます。 - 単独評価の指標:
mAP。evaluate アクションはステータス KPI にmAPを書き込み、コンソールの表には CMC ランクを表示します。単独の evaluate ステータスにcmc_rank_1が含まれることは想定しないでください。
アクションごとのデータセット要件
| アクション | Spec キー | ソース | ファイル | リスト? |
|---|---|---|---|---|
| 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 | いいえ |
一般的な Spec オーバーライド
すべてのアクションでデータソースのオーバーライドが必須です。エージェントは、上記の「アクションごとのデータセット要件」テーブルに基づいてデータソースのパスを構築し、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 の値または spec_params マッピングとして指定してください。ローカル Docker の skill_info.yaml ではファイル出力として宣言しないでください。ランナーが出力の事前作成時にファイルとフォルダーを区別できるようになるまで、この指定を行わないでください。
評価データセット
必須です。評価では、検索ベースの指標(CMC、mAP)にテストデータセットとクエリデータセットが必要です。
重要なパラメーター
- dataset.num_classes:ID の数です。デフォルトは 751 です。トレーニングデータ内の一意な ID の数と一致する必要があります。
- model.backbone:デフォルトは resnet_50 です。
- optim.base_lr:基本学習率です。デフォルトは 3.5e-4 です。
- dataset.batch_size:GPU あたりのバッチサイズです。デフォルトは 64 です。Re-ID では、トリプレット/コントラストサンプリングの精度向上に大きなバッチサイズが有効です。
- dataset.num_instances:バッチ内の ID ごとのインスタンス数です。メトリックラーニングのサンプリング戦略を制御します。
マルチ GPU/マルチノード
起動方法: Lightning による管理(単一の python プロセスから Lightning がワーカーを起動)。
| Spec キー | 説明 | デフォルト |
|---|---|---|
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 以上の VRAM が必要です。Re-ID モデルは比較的軽量ですが、メトリックラーニングには大きなバッチサイズが有効です。
エラーパターン
num_classes の不一致: dataset.num_classes が、トレーニングセット内の一意な ID フォルダーの数と一致していることを確認してください。
無効なトリプレットバッチ形状: dataset.batch_size は、各ミニバッチをハードサンプルマイニング用に再構成できるよう、dataset.num_instances と互換性がある必要があります。ローカルで AutoML のスモークテストを実行する場合は、dataset.batch_size を dataset.num_instances: 4 に対して 16 など既知の有効な倍数に固定し、制約のないバッチサイズを使うのではなく、train.optim.base_lr を調整してください。
クエリ/ギャラリーの不一致: クエリデータセットとテスト(ギャラリー)データセットでは、同じ ID 名前空間を使用する必要があります。
チェックポイントを読み込む側で発生する PyTorch 2.6 のチェックポイント読み込み失敗: 現在の Re-ID
チェックポイントには OmegaConf コンテナーが含まれています。同じ信頼できる TAO のトレーニング/AutoML ワークフローで生成されたチェックポイントについては、Lightning/PyTorch がチェックポイント全体を読み込めるよう、下流の resume、evaluate、inference、export ジョブの環境変数に
TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 を設定してください。この環境変数は、信頼できないチェックポイントには使用しないでください。
AutoML のメトリクス抽出: Re-ID のトレーニングステータスファイルには、cmc_rank_1、cmc_rank_5、cmc_rank_10、mAP などの検索 KPI と、トレーニング損失が報告されます。デフォルトの AutoML トレーニング起動では、direction: maximize を指定して cmc_rank_1 を最適化する必要があります。このモデルのメトリクスに val_loss を使用しないでください。
チェックポイントの受け渡し: 最適な AutoML 子ジョブの results_dir/train/ フォルダーにあるチェックポイントリゾルバーを使用して、アクションに適した model_epoch_*.pth チェックポイントを選択してください。Re-ID は reid_model_latest.pth も書き込みますが、これは最新のチェックポイントを指すシンボリックリンクであり、呼び出し元から最新のチェックポイントが明示的に要求された場合にのみ使用してください。下流のアクションでも、同じデータセット ID 数とクエリ/ギャラリーのアーカイブを維持してください。
デフォルト仕様の生成: パッケージに含まれる default_specs CLI アクションでは、results_dir の指定に通常の -e <spec.yaml> 実験ファイルを使用しません。
re_identification default_specs results_dir=/workspace/run/results/default_specs
のような Hydra オーバーライドを指定して実行してください。-e のみを渡すと cfg.results_dir が未設定のままとなり、
MissingMandatoryValue: results_dir
が発生します。
Spec Param/親モデルの推論
モデル固有の推論マッピングは、config.json ではなくこの MD ファイルに記述します。生成されたランナーはこのセクションを読み込み、create_job() の前に SDK ヘルパーを使用してマッピングを適用します。これは、以前のマイクロサービスにおける infer_params.py のフローを踏襲しています。
TAO Core re_identification.config.json の推論マッピング:
| アクション | Spec フィールド | 推論関数 | 意味 |
|---|---|---|---|
| 評価 | encryption_key | key | 暗号化キー |
| 評価 | evaluate.checkpoint | parent_model | 親ジョブの結果フォルダーから推論されるモデルファイル |
| 評価 | evaluate.output_cmc_curve_plot | create_evaluate_cmc_plot_reid | ReID CMC プロットのパス |
| 評価 | evaluate.output_sampled_matches_plot | create_evaluate_matches_plot_reid | ReID サンプリング済みマッチプロットのパス |
| 評価 | results_dir | output_dir | 現在のジョブの結果ディレクトリ |
| エクスポート | encryption_key | key | 暗号化キー |
| エクスポート | export.checkpoint | parent_model | 親ジョブの結果フォルダーから推論されるモデルファイル |
| エクスポート | export.onnx_file | create_onnx_file | 出力 ONNX パス |
| エクスポート | results_dir | output_dir | 現在のジョブの結果ディレクトリ |
| 推論 | encryption_key | key | 暗号化キー |
| 推論 | inference.checkpoint | parent_model | 親ジョブの結果フォルダーから推論されるモデルファイル |
| 推論 | inference.output_file | create_inference_result_file_reid | ReID 推論 JSON のパス |
| 推論 | results_dir | output_dir | 現在のジョブの結果ディレクトリ |
| トレーニング | encryption_key | key | 暗号化キー |
| トレーニング | 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 の場合、上流の train/export/AutoML 子ジョブ ID を parent_job_id として渡します。SDK は親の結果フォルダーを一覧表示し、チェックポイントのアーティファクトをフィルタリングして、選択されたモデルファイルまたはフォルダーを返します。これらのマッピングを config.json に戻したり、チェックポイントのパスを推測するよう生成されたランナースクリプトにパッチを適用したりしないでください。