tao-train-mask-auto-encoder

MAE

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

MAE（Masked Autoencoder）は、自己教師ありの事前学習とファインチューニングに使用します。ランダムなパッチをマスクし、それらを再構成することで、視覚表現を学習します。事前学習とファインチューニングのステージをサポートしています。

ファインチューニング時に事前学習済み MAE の重みを使用するには、train.pretrained_model_path を設定してください。

TAO Deploy の TensorRT アクション（gen_trt_engine）については、まず references/tao-deploy-mask-auto-encoder.md を読んでください。Deploy の仕様テンプレートは、このスキルの references/ フォルダーにあり、接頭辞は spec_template_deploy_*.yaml です。

親の PyTorch mae CLI は、train、evaluate、inference、export をサポートしています。 TensorRT エンジンはモデルスキルではなく、Deploy ワークフローを通じてビルドしてください。

データクラススキーマ

生成された 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 以外のアクションと Deploy フローは、このモデルスキルにとどまります。実行ごとの automl_policy の上書きによって、モデルのメタデータが変わることはありません。

トレーニング要件

• データセットタイプ： image_classification
• 形式： ssl
• 受け付けるデータセットの用途： training、evaluation、testing
• 事前学習の監視メトリクス： train_loss（最小化）。
• AutoML メトリクスの要件： ファインチューニングでは、最大化するメトリクスとして ACC_all を使用してください。ファインチューニングの検証とパッケージ化された evaluate アクションでは、train_loss は出力されず、ACC_all（および val_loss）が出力されます。train.stage: finetune の場合は、試行と最終チェックポイントの選択に ACC_all を使用してください。

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

アクション	スペックキー	ソース	ファイル	リスト?
train	dataset.train_data_sources	train_datasets	images_train.tar.gz	いいえ
train	dataset.val_data_sources	eval_dataset	images_val.tar.gz	いいえ
evaluate	dataset.val_data_sources	eval_dataset	images_val.tar.gz	いいえ
inference	dataset.test_data_sources	inference_dataset	images_test.tar.gz	いいえ

SDK/アプリのジョブ入力では、images_*.tar.gz アーカイブがアクション入力としてアップロードされます。ホストにマウントされたデータに対してローカルの Docker を直接実行する場合は、まずアーカイブを展開し、dataset.train_data_sources、dataset.val_data_sources、dataset.test_data_sources を展開後の images_train、images_val、images_test フォルダーに向けてください。ローカルの tar パスを直接 MAE CLI に渡すと、ローカルのデータローダーがそのアーカイブパスを展開しないため、サンプル数ゼロのデータローダーが生成されることがあります。

一般的なスペックのオーバーライド

データソースのオーバーライドはすべてのアクションで必須です。エージェントは上記のアクションごとのデータセット要件表からデータソースパスを構築し、spec_overrides に含めなければなりません (MUST)。

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

train (必須のデータソース):

{
    "dataset.train_data_sources": f"{S3_TRAIN}/images_train.tar.gz",
    "dataset.val_data_sources": f"{S3_EVAL}/images_val.tar.gz",
    "train.num_epochs": 10,
    "train.optim.lr": 2e-4,
}

evaluate (必須のデータソース):

{
    "dataset.val_data_sources": f"{S3_EVAL}/images_val.tar.gz",
    "evaluate.checkpoint": "<selected train/AutoML checkpoint>",
    "train.stage": "finetune",
}

inference (必須のデータソース):

{
    "dataset.test_data_sources": f"{S3_EVAL}/images_test.tar.gz",
    "inference.checkpoint": "<selected train/AutoML checkpoint>",
    "train.stage": "finetune",
}

Eval データセット

任意です。事前学習には eval データは必要ありません。ファインチューニングではオプションで val セットを使用します。

重要なパラメーター

• train.stage: トレーニング段階。オプション: pretrain、finetune。pretrain はマスキングを通じて表現を学習します。finetune は分類ヘッドを追加します。
• model.arch: アーキテクチャ。デフォルトは convnextv2_base。ローカルのスモーク AutoML では、vit_tiny_patch16 のようなサポートされていない名前ではなく convnextv2_atto を使用してください。サポートされるファミリーには vit_base_patch16 およびそれ以上の ViT、ConvNeXtV2 の atto/femto/pico/nano/tiny/base/large/huge、そして Hiera の tiny/small/base/large/huge が含まれます。
• model.num_classes: ファインチューニング用のクラス数。デフォルトは 1000 (ImageNet)。finetune 段階でのみ関連します。
• model.mask_ratio: 事前学習中にマスクするパッチの割合。通常は 0.75。
• model.norm_pix_loss: 再構成損失でピクセル値を正規化するかどうか。
• dataset.augmentation.input_size: ConvNeXtV2 MAE ではローカルのスモークプロファイルを 224 のままにしてください。 112 に減らすと、MAE のマスクグリッドが特徴マップの次元と互換性がなくなる可能性があります。
• MAE は dataset.workers スペックフィールドを公開していません。スモークテストの オーバーライドに追加しないでください。Hydra はトレーニング前に未知の dataset キーを拒否します。
• train.optim.lr: 学習率。デフォルトは 2e-4。
• dataset.augmentation: ファインチューニング用の mixup、cutmix を含む拡張設定。

マルチ GPU / マルチノード

起動方法: Lightning 管理 (単一の python プロセスで、Lightning がワーカーを生成)。

スペックキー	説明	デフォルト
train.num_gpus	GPU の数	1
train.gpu_ids	GPU デバイスインデックス	[0]
train.num_nodes	ノード数	1
train.distributed_strategy	ddp または fsdp	ddp

• ddp は find_unused_parameters=True を使用します
• fsdp は FP16 を強制します
• 事前学習にはマルチ GPU を強く推奨します (大きなバッチサイズが必要)

マルチノード環境変数 (オーケストレーターによって設定): WORLD_SIZE、NODE_RANK、MASTER_ADDR、MASTER_PORT、NUM_GPU_PER_NODE。

ハードウェア

最低 2 GPU、推奨 8 GPU。GPU あたり 24GB 以上 (A100 推奨) の VRAM。MAE の事前学習は多くの GPU にまたがる大きなバッチサイズの恩恵を受けます。ファインチューニングはリソース要件がより控えめです。

エラーパターン

ステージの不一致：train.stage が意図（pretrain または finetune）に合っていることを確認してください。pretrained_model_path を指定せずにファインチューニングすると、スクラッチから学習します。

pretrain チェックポイントを使用した推論：MAE の predict データローダーは、train.stage: pretrain の場合に NotImplementedError を発生させます。推論と分類形式の評価には finetune チェックポイントを使用するか、pretrain のみの実行を train/evaluate/export に限定してください。

num_classes の不一致（finetune のみ）：ファインチューニング時は、model.num_classes がデータセットのクラス数と一致していることを確認してください。

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

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

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

アクション	Spec フィールド	推論関数	意味
evaluate	encryption_key	key	暗号化キー
evaluate	evaluate.checkpoint	parent_model	親ジョブの結果フォルダーから推定されるモデルファイル
evaluate	evaluate.trt_engine	parent_model	親ジョブの結果フォルダーから推定されるモデルファイル
evaluate	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 に戻したり、チェックポイントのパスを推測するよう生成済みランナースクリプトにパッチを適用したりしないでください。

SDK のリゾルバー以外でチェックポイントを解決する場合は、たとえば model_epoch_000_step_00099.pth のように、対象のエポック／ステップのアーティファクトを正確に選択してください。convnextv2_atto_latest.pth やその他の latest シンボリックリンクを使用するのは、latest が明示的に要求された場合に限ります。チェックポイントと ONNX／エンジンの形状が一致するよう、train.stage、model.arch、model.num_classes、およびエクスポート時の入力サイズを evaluate、inference、export、deploy の各 spec に引き継いでください。

デプロイ

• tao-deploy-mask-auto-encoder