tao-train-metric-learning-recognition

ML Recog

細粒度の視覚認識に向けたメトリック学習による認識です。検索ベースのマッチング（例：小売商品の認識）に使用する埋め込みを学習します。トリプレット損失／コントラスト損失を使用します。

事前学習済みバックボーンを使用するには、model.pretrained_model_path を設定します。

TAO Deploy の TensorRT アクション（gen_trt_engine、TensorRT の evaluate、TensorRT の inference）については、まず references/tao-deploy-metric-learning-recognition.md をお読みください。デプロイ仕様テンプレートは、このスキルの references/ フォルダーにあり、プレフィックスは spec_template_deploy_*.yaml です。

データクラススキーマ

生成された TAO Core スキーマは schemas/<action>.schema.json にパッケージ化されており、利用可能なアクションは schemas/manifest.json に記載されています。生成された各スキーマから、スキーマのトップレベルの default フィールドをもとに references/spec_template_<action>.yaml も生成されます。AutoML の有効化は、モデルレイヤーの references/skill_info.yaml で automl_enabled を使って宣言します。AutoML を実行するには、schemas/train.schema.json と references/spec_template_train.yaml が存在し、解析可能である必要があります。automl_default_parameters、automl_disabled_parameters、デフォルト値、最小値／最大値の境界、列挙値、オプションの重み、数式条件、依存関係、よく使われるパラメーターには、パッケージ化された train スキーマを使用してください。実行時に ~/tao-core が存在することを前提にしないでください。メンテナーは、スキルバンクをパッケージ化する前にスキーマとテンプレートを再生成してください。

Train アクションのポリシー

このモデルでは、モデルレイヤーで AutoML が有効になっています。train ステージのリクエストを処理する前に、references/skill_info.yaml を読み、明示的な automl_policy の値、またはユーザーのワークフローリクエストから実行時の上書きを判断してください。「AutoML をオフにする」「AutoML を無効にする」「HPO なしで」「通常のトレーニング」といった表現は、この実行に限り automl_policy: off として扱ってください。それ以外の場合は、デフォルトの auto を使用します。automl_policy: auto、automl_enabled: true であり、schemas/train.schema.json と references/spec_template_train.yaml の両方がパッケージ化されている場合は、デフォルトでこのモデルの skill_dir を指定して、train アクションを tao-skill-bank:tao-run-automl 経由で実行してください。データセット、仕様、出力ディレクトリ、GPU／プラットフォーム設定、親チェックポイント、automl_policy に関するワークフロー／アプリケーションの上書きは維持してください。直接モデルをトレーニングするのは、automl_policy: off の場合、またはパッケージ化された train スキーマ／テンプレートが存在しない場合に限ります。スキーマが存在しない場合は、このモデルでは AutoML が有効ですが、スキーマが生成されるまで実行できないことを報告してください。

evaluate、inference、export、デプロイフローなどの train 以外のアクションは、このモデルのスキル内で引き続き処理します。実行ごとの automl_policy 上書きによって、モデルのメタデータが変更されることはありません。

トレーニング要件

• データセットタイプ： ml_recog
• 形式： デフォルト
• モニタリング指標： val Precision at Rank 1

アクションごとのデータセット要件

アクション	Spec Key	ソース	ファイル	リスト?
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	No
gen_trt_engine	gen_trt_engine.tensorrt.calibration.cal_image_dir	calibration_dataset	metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes/test.tar.gz	Yes
inference	dataset.val_dataset	train_datasets	reference: metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/reference.tar.gz, query:	No
inference	inference.input_path	train_datasets	metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/test.tar.gz	No
train	dataset.train_dataset	train_datasets	metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes/train.tar.gz	No
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	No

一般的な Spec オーバーライド

データソースのオーバーライドはすべてのアクションで必須です。エージェントは、上記の「アクションごとのデータセット要件」表からデータソースのパスを作成し、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"},
}

gen_trt_engine（必須データソース）：

{
    "gen_trt_engine.tensorrt.data_type": "INT8",
    "gen_trt_engine.tensorrt.calibration.cal_image_dir": [f"{S3_TRAIN}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes/test.tar.gz"],
}

evaluate（必須データソース）：

{
    "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（必須データソース）：

{
    "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 です。
• dataset.num_instance：バッチ内の ID ごとのインスタンス数（P/K サンプリング）です。デフォルトは 4 です。同じクラスの画像を何枚まとめて処理するかを指定します。
• 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：学習率の減衰ステップです。デフォルトは gamma=0.1 の [40, 70] です。
• dataset.train_dataset：クラスフォルダーに整理された学習画像のパスです。
• dataset.val_dataset：検索評価用の ImageNet 形式ディレクトリを指す 'reference' と 'query' キーを持つ辞書です。

マルチ GPU／マルチノード

起動方法： Lightning による管理（単一の python プロセスで、Lightning がワーカーを起動します）。

Spec Key	説明	デフォルト
train.num_gpus	GPU の数	1
train.gpu_ids	GPU デバイスのインデックス	[0]

• ストラテジー：auto（Lightning が最適なストラテジーを自動的に選択します）
• 明示的な num_nodes または distributed_strategy の設定はありません。単一ノード向けです。

ハードウェア

GPU は最低 1 基、推奨は 2 基です。GPU ごとに 16GB 以上の VRAM が必要です。メトリック学習では、より良いトリプレットサンプリングのために大きなバッチサイズが有効ですが、それ以外のメモリ使用量は中程度です。

エラーパターン

reference/query の不一致：評価時に、reference データセットと query データセットで互換性のあるクラス名前空間が使用されていることを確認してください。

仕様パラメーター／親モデルの推論

モデル固有の推論マッピングは、config.json ではなくこの MD ファイルに記述します。生成されたランナーはこのセクションを読み取り、create_job() の前に SDK ヘルパーを使用してマッピングを適用してください。これは、以前のマイクロサービスの infer_params.py フローに相当します。

TAO Core の ml_recog.config.json からの推論マッピング：

アクション	仕様フィールド	推論関数	意味
評価	evaluate.checkpoint	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
評価	evaluate.trt_engine	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
評価	results_dir	output_dir	現在のジョブの結果ディレクトリ
エクスポート	export.checkpoint	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
エクスポート	export.onnx_file	create_onnx_file	出力 ONNX パス
エクスポート	results_dir	output_dir	現在のジョブの結果ディレクトリ
gen_trt_engine	gen_trt_engine.onnx_file	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
gen_trt_engine	gen_trt_engine.tensorrt.calibration.cal_cache_file	create_cal_cache	キャリブレーションキャッシュのパス
gen_trt_engine	gen_trt_engine.trt_engine	create_engine_file	出力 TensorRT エンジンのパス
gen_trt_engine	results_dir	output_dir	現在のジョブの結果ディレクトリ
推論	inference.checkpoint	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
推論	inference.trt_engine	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 では、上流の train／export／AutoML 子ジョブ ID を parent_job_id として渡します。SDK は親の結果フォルダーを一覧表示し、チェックポイントのアーティファクトを絞り込み、選択されたモデルファイルまたはフォルダーを返します。これらのマッピングを config.json に戻したり、チェックポイントのパスを推測するよう生成されたランナースクリプトにパッチを適用したりしないでください。

デプロイ

• tao-deploy-metric-learning-recognition — TAO Deploy を使用した TensorRT エンジン生成、TensorRT 評価、TensorRT 推論のための MLRecog デプロイワークフロー。