tao-train-mask-auto-encoder

MAE

自己教師あり事前学習およびファインチューニング用の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というプレフィックスが付いています。

データクラススキーマ

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

トレーニング要件

• データセットタイプ： image_classification
• 形式： ssl
• 受け付けるデータセットの用途： トレーニング、評価、テスト
• 監視メトリック： train_loss

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

アクション	仕様キー	ソース	ファイル	リスト？
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	いいえ

代表的な仕様の上書き

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

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",
}

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

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

評価データセット

任意です。事前学習では評価データは不要です。ファインチューニングでは、任意で検証セットを使用します。

重要なパラメーター

• train.stage：トレーニングのステージです。選択肢は pretrain、finetune です。Pretrain ではマスキングによって表現を学習します。Finetune では分類ヘッドを追加します。
• model.arch：アーキテクチャです。デフォルトは convnextv2_base です。ConvNeXt、Hiera、ViT の各種モデルなど、幅広い選択肢があります。
• model.num_classes：ファインチューニング用のクラス数です。デフォルトは 1000（ImageNet）です。finetune ステージでのみ使用します。
• model.mask_ratio：事前学習中にマスクするパッチの割合です。通常は 0.75 です。
• model.norm_pix_loss：再構成損失でピクセル値を正規化するかどうかを指定します。
• 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 を使用した大きなバッチサイズが効果的です。ファインチューニングでは、必要なリソースは比較的少なくなります。

エラーのパターン

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

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

Spec Param / 親モデルの推論

モデル固有の推論マッピングは 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	現在のジョブの結果ディレクトリ
gen_trt_engine	encryption_key	key	暗号化キー
gen_trt_engine	gen_trt_engine.onnx_file	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
gen_trt_engine	gen_trt_engine.trt_engine	create_engine_file	出力 TensorRT エンジンパス
gen_trt_engine	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-mask-auto-encoder — TAO Deploy を使用して TensorRT エンジンを生成する MAE デプロイワークフロー。