tao-train-metric-learning-recognition

ML Recog

単独でインストールする場合 このセッションが TAO skill bank プラグインによって初期化されていない場合は、最初に tao-setup スキルを実行してください（ホストの事前確認、認証情報、スキル間の検出）。

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

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

TAO Deploy の TensorRT アクション（gen_trt_engine、TensorRT の evaluate、TensorRT の inference）については、最初に references/tao-deploy-metric-learning-recognition.md を読んでください。Deploy 仕様テンプレートは、このスキルの 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 があることを前提にしないでください。メンテナーはスキルバンクをパッケージ化する前に、スキーマ／テンプレートを再生成してください。

Train アクションのポリシー

このモデルでは、モデルレイヤーで AutoML が有効になっています。train ステージのリクエストに対応する前に、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 にこのモデルのパスを指定して、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
• AutoML 指標の条件： トレーニング中に出力される val Precision at Rank 1 を使用し、最大化してください。test Precision at Rank 1 は、選択したチェックポイントの単独検証にのみ使用してください。
• 単独評価指標： test Precision at Rank 1。AutoML の推奨ランキングには、トレーニング時の val Precision at Rank 1 KPI を使用し、評価用の参照／クエリ分割で選択したチェックポイントを検証する場合にのみ 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 の上書き

データソースの上書きはすべてのアクションで必須です。エージェントは、上記の「アクションごとのデータセット要件」表からデータソースのパスを構成し、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：バッチ内のIDごとのインスタンス数（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：学習率の減衰ステップです。デフォルトはgamma=0.1で[40, 70]です。
• dataset.train_dataset：クラス別フォルダーに整理された学習画像へのパスです。
• dataset.val_dataset：検索評価用のImageNet形式ディレクトリを指す'reference'キーと'query'キーを含む辞書です。

マルチ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以上のVRAMが必要です。メトリック学習では、トリプレットサンプリングの精度向上に大きなバッチサイズが有効ですが、それ以外のメモリ使用量は中程度です。

エラーパターン

参照データセットとクエリデータセットの不一致：評価用に、参照データセットとクエリデータセットのクラス名前空間に互換性があることを確認してください。

チェックポイント関連のアクションで発生する PyTorch 2.6 のチェックポイント読み込みエラー：現在の TAO ML-Recog チェックポイントには、OmegaConf オブジェクトが含まれている場合があります。同じ信頼できる TAO の train／AutoML ワークフローで生成されたチェックポイントの場合は、後続の evaluate、inference、export、 または resume/retrain ジョブの環境変数に TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 を設定すると、Lightning でチェックポイント全体を読み込めるようになります。この環境変数は、信頼できないチェックポイントには使用しないでください。

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

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

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

アクション	Spec フィールド	推論関数	意味
evaluate	evaluate.checkpoint	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
evaluate	results_dir	output_dir	現在のジョブの結果ディレクトリ
export	export.checkpoint	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
export	export.onnx_file	create_onnx_file	出力 ONNX パス
export	results_dir	output_dir	現在のジョブの結果ディレクトリ
inference	inference.checkpoint	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
inference	results_dir	output_dir	現在のジョブの結果ディレクトリ
train	model.pretrained_model_path	ptm_if_no_resume_model	再開用チェックポイントがない場合の PTM
train	results_dir	output_dir	現在のジョブの結果ディレクトリ
train	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