tao-train-nvdinov2

NVDINOv2

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

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

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

TAO Deploy の TensorRT アクション（gen_trt_engine）については、最初に references/tao-deploy-nvdinov2.md をお読みください。Deploy 仕様テンプレートは、このスキルの 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/<action>.schema.json と references/spec_template_<action>.yaml が存在し、解析できる必要があります。automl_default_parameters、automl_disabled_parameters、デフォルト値、最小値・最大値の範囲、列挙値、オプションの重み、数学的条件、依存関係、よく使われるパラメーターには、パッケージ化された選択アクションのスキーマを使用してください。実行時に ~/tao-core が存在するとは想定しないでください。メンテナーはスキルバンクをパッケージ化する前に、スキーマとテンプレートを再生成します。

学習アクションのポリシー

このモデルではモデルレイヤーで 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 経由で学習アクションを実行してください。データセット、仕様、出力ディレクトリ、GPU／プラットフォーム設定、親チェックポイント、automl_policy に対するワークフローやアプリケーションのオーバーライドは維持してください。automl_policy: off の場合、またはパッケージ化された学習スキーマ／テンプレートがない場合に限り、モデルを直接学習してください。スキーマがない場合は、AutoML は有効ですが、スキーマが生成されるまではこのモデルで実行できないことを報告してください。

inference、export、デプロイフローなど、学習以外のアクションはこのモデルスキル内にとどまります。実行ごとの automl_policy オーバーライドによって、モデルのメタデータが変わることはありません。

学習要件

• データセットタイプ： image_classification
• 形式： ssl
• モニタリング指標： train_loss
• AutoML 指標の契約： 学習アクションが出力する train_loss を使用し、 最小化してください。このモデルにはパッケージ化された評価アクションがないため、 評価 KPI を独自に設定しないでください。選択した student_epoch_* チェックポイントは 推論を通じて検証してください。

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

アクション	仕様キー	ソース	ソースアーティファクト → 実行時の値	リスト？
inference	dataset.test_dataset.images_dir	inference_dataset	images_test.tar.gz → 展開済み画像フォルダー	いいえ
train	dataset.train_dataset.images_dir	train_datasets	images_train.tar.gz → 展開済み画像フォルダー	いいえ

references/skill_info.yaml で runtime: extracted_folder が宣言されている場合、管理対象ランナーはこれらのデータセットをアーカイブとして取得できます。アーカイブを展開し、生成されたディレクトリを images_dir に配置する必要があります。Docker または TAO CLI を直接実行する場合は、起動前にアーカイブを展開してください。images_dir フィールドに .tar、.tar.gz、.tgz、.zip を指定しないでください。

典型的な Spec Overrides

データソースのオーバーライドはトレーニングと推論で必須です。エージェントは、上記の Per-Action Dataset Requirements テーブルに基づいてデータソースのパスを構築し、spec_overrides に含める必要があります。

TRAIN_IMAGES_DIR = "/workspace/data/extracted/train/images_train"
TEST_IMAGES_DIR = "/workspace/data/extracted/test/images_test"

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

{
    "train.num_gpus": 1,
    "train.num_epochs": 10,
    "train.checkpoint_interval": 10,
    "dataset.train_dataset.images_dir": TRAIN_IMAGES_DIR,
}

ローカル AutoML 検証／スモーク実行： ローカルの Docker でベイズ最適化の起動、メトリックの選択、最良モデルの選択、チェックポイントの永続化を確認する場合は、この形式を使用します。実行の代表性を保ちながら、はるかに時間のかかる ViT-Large のデフォルト設定を避けられます。

{
    "wandb.enable": False,
    "model.backbone.teacher_type": "vit_s",
    "model.backbone.student_type": "vit_s",
    "model.backbone.img_size": 224,
    "dataset.batch_size": 8,
    "dataset.workers": 2,
    "train.num_epochs": 1,
    "train.checkpoint_interval": 1,
    "train.num_prototypes": 1024,
    "train.precision": "32-true",
    "train.use_custom_attention": False,
    "train.num_gpus": 1,
    "dataset.train_dataset.images_dir": TRAIN_IMAGES_DIR,
}

export（必須のチェックポイント引き渡し）：

{
    "export.checkpoint": "<selected train/AutoML student_epoch_* checkpoint>",
    "export.onnx_file": "/path/to/results/nvdinov2.onnx",
}

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

{
    "inference.checkpoint": "<selected train/AutoML student_epoch_* checkpoint>",
    "model.backbone.teacher_type": "<same value used for train>",
    "model.backbone.student_type": "<same value used for train>",
    "model.backbone.img_size": "<same value used for train>",
    "train.use_custom_attention": "<same value used for train>",
    # NVDINOv2 enables persistent_workers in its prediction DataLoader.
    # Keep this positive; dataset.workers=0 makes inference fail before loading data.
    "dataset.workers": 2,
    "dataset.test_dataset.images_dir": TEST_IMAGES_DIR,
}

評価データセット

任意です。SSL トレーニングではラベルを使用しません。評価は下流タスクごとに異なります。

重要なパラメーター

• model.backbone.teacher_type：教師 ViT のバリアントです。デフォルトは vit_l（ViT-Large）です。
• model.backbone.student_type：生徒 ViT のバリアントです。デフォルトは vit_l です。通常は教師と同じものを指定します。
• 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 Key	説明	デフォルト
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 ごとに 40GB 以上の VRAM が必要です（A100 を推奨）。ViT-Large の教師と生徒を使用する SSL は、非常に多くのメモリーを必要とします。A100 40GB 以上の GPU が必要です。マルチ GPU を強く推奨します。

エラーパターン

画像が見つからない／アーカイブのパスが拒否される：dataset.*.images_dir には、アーカイブではなく画像を含む展開済みディレクトリを指定する必要があります。.tar、.tar.gz、.tgz、.zip の入力は、先に展開してください。アーカイブを使用するマネージドソースでは、runtime: extracted_folder によって生成されたディレクトリを使用する必要があります。

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

推論チェックポイントに予期しない Lightning のキーが含まれる：下流の inference には、nvdinov2_model_latest.pth ではなく、選択した AutoML 実行の student_epoch_*.pth チェックポイントを渡してください。latest ファイルはトレーニングチェックポイントであり、 推論ローダーは state_dict、オプティマイザーの状態、スケジューラーの状態などの予期しないキーを報告します。

推論が persistent_workers option needs num_workers > 0 で失敗する： dataset.workers を正の値に設定してください（最小構成のローカル実行では 2 を使用します）。 NVDINOv2 の予測 DataLoader は永続ワーカーを有効にするため、 汎用のゼロワーカーのスモークテスト用オーバーライドでは実行できません。

エクスポートチェックポイントに予期しない Lightning のキーが含まれる：エクスポートでも、選択した student_epoch_*.pth チェックポイントを使用します。完全な model_epoch_*.pth チェックポイントは、train.resume_training_checkpoint_path による再開／再トレーニングにのみ使用してください。

PyT 推論に渡された TensorRT エンジン：パッケージ版の PyT nvdinov2 inference 実装が読み込めるモデルパスは .pth または .tlt のみです。TAO Deploy の gen_trt_engine は下流の利用側向けに TensorRT エンジンをビルドしますが、PyT 推論アクションはそのエンジンでは実行されません。

個別の distill アクションを利用できない：現在の TAO PyT CLI が NvDINOv2 向けに公開しているのは、 export、inference、train、default_specs です。独立した nvdinov2 distill アクションを起動したり、利用可能なものとして案内したりしないでください。

AutoML のメトリクスが見つからない：TAO のステータス KPI では、トレーニングの最終スカラー値は train_loss として報告されます。AutoML の選択には、最小化方向で train_loss を使用してください。 一部の Lightning の進捗行では、同じスカラー値が train_loss_epoch として表示されることもあります。これはフォールバック用の別名としてのみ扱い、主要な モニタリングメトリクスにはしないでください。

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

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

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

モデル固有の引き渡しマッピング：

アクション	Spec フィールド	推論関数	意味
export	encryption_key	key	暗号化キー
export	export.checkpoint	parent_model	親ジョブの結果フォルダーにある、選択済みの student_epoch_*.pth チェックポイント
export	export.onnx_file	create_onnx_file	出力 ONNX パス
export	results_dir	output_dir	現在のジョブの結果ディレクトリ
inference	encryption_key	key	暗号化キー
inference	inference.checkpoint	parent_model	親ジョブの結果フォルダーにある、選択済みの student_epoch_*.pth チェックポイント
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	現在のジョブの結果フォルダーにある、選択済みの完全な model_epoch_*.pth トレーニングチェックポイント

parent_model または parent_model_folder には、上流の train/export/AutoML 子ジョブ ID を parent_job_id として渡してください。SDK は親の結果フォルダーを一覧表示し、チェックポイントのアーティファクトを絞り込んで、選択されたモデルファイルまたはフォルダーを返します。これらのマッピングを config.json に戻したり、チェックポイントパスを推測するように生成済みのランナースクリプトを修正したりしないでください。

Deployment

• tao-deploy-nvdinov2