tao-train-nvdinov2

NVDINOv2

NVDINOv2 は、自己教師ありの視覚表現学習向けです。ラベルなしで自己蒸留（教師・生徒方式）によりビジョントランスフォーマーを学習します。汎用的な視覚特徴を生成します。

事前学習済みの ViT 重みには、train.pretrained_model_path を設定してください。

TAO Deploy の TensorRT アクション（gen_trt_engine）については、先に references/tao-deploy-nvdinov2.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/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 を指定して tao-skill-bank:tao-run-automl 経由で train アクションを実行します。データセット、仕様、出力ディレクトリ、GPU／プラットフォーム設定、親チェックポイント、automl_policy に対するワークフロー／アプリケーションの上書きは保持してください。直接モデル学習を行うのは、automl_policy: off の場合、またはパッケージ化された train スキーマ／テンプレートがない場合に限ります。スキーマがない場合は、AutoML が有効ではあるものの、スキーマが生成されるまではこのモデルで実行できないことを報告してください。

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

学習要件

• データセットタイプ： image_classification
• フォーマット： ssl
• 監視メトリック： train_loss

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

アクション	仕様キー	ソース	ファイル	リストか？
distill	dataset.train_dataset.images_dir	train_datasets	images_train.tar.gz	いいえ
inference	dataset.test_dataset.images_dir	inference_dataset	images_test.tar.gz	いいえ
train	dataset.train_dataset.images_dir	train_datasets	images_train.tar.gz	いいえ

一般的な仕様の上書き

データソースの上書きはすべてのアクションで必須です。エージェントは、上記の「アクションごとのデータセット要件」テーブルに基づいてデータソースのパスを構築し、spec_overrides に含める必要があります。

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

train（必須のデータソース）：

{
    "train.num_gpus": 1,
    "train.num_epochs": 10,
    "train.checkpoint_interval": 10,
    "dataset.train_dataset.images_dir": f"{S3_TRAIN}/images_train.tar.gz",
}

distill（必須のデータソース）：

{
    "dataset.train_dataset.images_dir": f"{S3_TRAIN}/images_train.tar.gz",
}

inference（必須のデータソース）：

{
    "dataset.test_dataset.images_dir": f"{S3_EVAL}/images_test.tar.gz",
}

評価データセット

任意です。SSL 学習ではラベルを使用しません。評価は下流タスクに応じて行います。

重要なパラメーター

• model.backbone.teacher_type：Teacher ViT のバリアントです。デフォルトは vit_l（ViT-Large）です。
• model.backbone.student_type：Student ViT のバリアントです。デフォルトは vit_l です。通常は Teacher と同じものを使用します。
• model.backbone.img_size：入力画像サイズです。デフォルトは 518 です。解像度を高くすると、より優れた特徴が得られますが、メモリ使用量が増加します。
• model.backbone.patch_size：ViT のパッチサイズです。デフォルトは 14 です。
• dataset.batch_size：GPU あたりのバッチサイズです。デフォルトは 4 です。教師と学生の両方で順伝播を行うため、SSL トレーニングはメモリを大量に消費します。
• train.layerwise_decay：レイヤー単位の学習率減衰です。ViT のファインチューニングで重要です。
• train.clip_grad_norm：勾配クリッピングです。安定した SSL トレーニングで重要です。

マルチ GPU／マルチノード

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

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

• 戦略：auto（Lightning が最適な戦略を自動的に選択します）
• sync_batchnorm は常に有効です。教師と学生のフレームワークによる SSL トレーニングでは非常に重要です。
• 有意義な SSL トレーニングには、マルチ GPU（4～8 GPU）を強く推奨します。

マルチノード環境変数（オーケストレーターが設定）：WORLD_SIZE、NODE_RANK、MASTER_ADDR、MASTER_PORT、NUM_GPU_PER_NODE。

ハードウェア

最低 4 GPU、推奨 8 GPU です。GPU あたり 40 GB 以上（A100 を推奨）の VRAM が必要です。ViT-Large の教師と学生による SSL は、非常に多くのメモリを消費します。A100 40 GB 以上の GPU が必要です。マルチ GPU を強く推奨します。

エラーパターン

CUDA のメモリ不足：img_size=518 の ViT-Large の教師と学生には、40 GB 以上の GPU メモリが必要です。batch_size または img_size を小さくするか、より小さい ViT バリアントを使用してください。

収束が遅い：SSL には多くのエポックが必要です。デフォルトの 10 は簡易テスト用です。本番環境での実行では通常、100 エポック以上を使用します。

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

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

TAO Core の nvdinov2.config.json に基づく推論マッピング：

アクション	Spec フィールド	推論関数	意味
distill	encryption_key	key	暗号化キー
distill	model.distill.pretrained_non_distill_pl_model_path	parent_model	親ジョブの結果フォルダーから推論されたモデルファイル
distill	results_dir	output_dir	現在のジョブの結果ディレクトリ
export	encryption_key	key	暗号化キー
export	export.checkpoint	parent_model	親ジョブの結果フォルダーから推論されたモデルファイル
export	export.onnx_file	create_onnx_file	出力 ONNX パス
export	results_dir	output_dir	現在のジョブの結果ディレクトリ
inference	encryption_key	key	暗号化キー
inference	inference.checkpoint	parent_model	親ジョブの結果フォルダーから推論されたモデルファイル
inference	inference.trt_engine	parent_model	親ジョブの結果フォルダーから推論されたモデルファイル
inference	results_dir	output_dir	現在のジョブの結果ディレクトリ
train	encryption_key	key	暗号化キー
train	results_dir	output_dir	現在のジョブの結果ディレクトリ
train	train.pretrained_model_path	ptm_if_no_resume_model	再開用チェックポイントが存在しない場合の PTM
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-nvdinov2 — TAO Deploy を使用して TensorRT エンジンを生成する NvDINOv2 のデプロイワークフローです。