mcpField
T
tao-train-mask-auto-label

tao-train-mask-auto-label

github.com/nvidiaデータ・データベース3.8k1.6k 回インストール今日更新
リポジトリ

tao-train-mask-auto-label について

MALは、点やボックスなどの少量のアノテーションからマスクを生成する、弱教師ありセグメンテーション向けのTAOモデルスキルです。学習・評価・推論に対応し、学習時にはAutoMLによるパラメーター探索も利用できます。

主な機能

  1. 01少量のアノテーションからマスクを生成
  2. 02ViT-MAEバックボーンを使用
  3. 03学習・評価・推論に対応
  4. 04学習時のAutoML探索

活用シーン

  1. 01点やボックスのラベルからセグメンテーションを学習
  2. 02検証データでmIoUを評価
  3. 03画像にセグメンテーション推論を実行

インストール

コマンドをコピーしてターミナルで実行します。

$npx skills add https://github.com/nvidia/skills --skill tao-train-mask-auto-label

インストール後、Claude Code で/tao-train-mask-auto-labelを入力して読み込まれたか確認します。

同梱ファイル

SKILL.md と一緒にインストールされる参考資料やスクリプト。

BENCHMARK.md
skill-card.md
skill.oms.sig
config/
skillspector-baseline.yaml
evals/
evals.json
references/
skill_info.yaml
spec_template_evaluate.yaml
spec_template_inference.yaml
spec_template_train.yaml
schemas/
evaluate.schema.json
inference.schema.json
manifest.json
train.schema.json

詳細

MAL

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

弱教師ありセグメンテーション用の MAL(Mask Auto-Label)です。最小限のアノテーション(例:ポイントまたはボックスアノテーション)からセグメンテーションマスクを生成します。ViT-MAE バックボーンを使用します。

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

クイックスタート(docker run)

Docker ネイティブで起動できるため、ホストに TAO SDK や Python は不要です。Docker/プラットフォーム用スキルのほうが、より厳密な環境固有のコマンド(非 root UID のマッピング、キャッシュのリダイレクト、リモートデーモンなど)を提供する場合は、そちらを使用してください。

TAO_PYT_IMAGE_DEFAULT=nvcr.io/nvidia/tao/tao-toolkit:7.3.0-pyt  # versions-key: images.tao_toolkit.pyt
TAO_PYT_IMAGE="${TAO_PYT_IMAGE:-$TAO_PYT_IMAGE_DEFAULT}"
RUN_ROOT="${RUN_ROOT:-$PWD}"
DOCKER_COMMON=(
  --rm --gpus all --shm-size=8g
  --shm-size=8g
  --ulimit memlock=-1
  --ulimit stack=67108864
  -v "$RUN_ROOT/data:/data:ro"
  -v "$RUN_ROOT/specs:/specs:ro"
  -v "$RUN_ROOT/results:/results"
)

トレーニング:

docker run "${DOCKER_COMMON[@]}" "$TAO_PYT_IMAGE" \
  mal train -e /specs/train.yaml

評価:

docker run "${DOCKER_COMMON[@]}" "$TAO_PYT_IMAGE" \
  mal evaluate -e /specs/evaluate.yaml

推論:

docker run "${DOCKER_COMMON[@]}" "$TAO_PYT_IMAGE" \
  mal inference -e /specs/inference.yaml

各アクションには -e で spec を指定します。results_dir は spec で設定するか、コマンドラインで上書きします。spec が参照する事前学習済み重みのディレクトリをマウントし、すべてのアクションでコンテナー内のパスを統一してください。

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 があることを前提にしないでください。メンテナーはスキルバンクのパッケージ化前に、スキーマとテンプレートを再生成してください。

トレーニングアクションのポリシー

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

evaluate、inference、export、デプロイフローなどの train 以外のアクションは、このモデルのスキル内で実行してください。実行ごとの automl_policy オーバーライドによって、モデルのメタデータは変更されません。

トレーニング要件

  • データセットの種類: セグメンテーション
  • フォーマット: デフォルト
  • モニタリング指標: mIoU
  • AutoML 指標の契約: evaluate アクションが出力する mIoU を使用し、最大化してください。記録された AutoML の目的関数を評価器が出力した mIoU と比較し、トレーニング損失で代用しないでください。
アクション別データセット要件
アクションSpec キーソースファイルリスト?
evaluatedataset.val_img_direval_datasetimages.tar.gzNo
evaluatedataset.val_ann_patheval_datasetannotations.jsonNo
inferenceinference.img_dirinference_datasetimages.tar.gzNo
inferenceinference.ann_pathinference_datasetannotations.jsonNo
traindataset.train_img_dirtrain_datasetsimages.tar.gzNo
traindataset.train_ann_pathtrain_datasetsannotations.jsonNo
traindataset.val_img_direval_datasetimages.tar.gzNo
traindataset.val_ann_patheval_datasetannotations.jsonNo
典型的な Spec オーバーライド

データソースのオーバーライドは すべてのアクションで必須 です — エージェントは上記のアクション別データセット要件の表からデータソースパスを構築し、それらを spec_overrides に含めなければなりません (MUST)。 MAL は、データソースが準備された後に、JSON の file_name エントリと一致する画像パスとともに、COCO 形式のアノテーション JSON を期待します。アーカイブのみの CSV/画像 データセットは、先にこの形式へ変換されない限り互換性がありません。

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

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

{
    "train.num_gpus": 1,
    "train.gpu_ids": [
        0
    ],
    "train.num_epochs": 5,
    "train.checkpoint_interval": 5,
    "train.validation_interval": 5,
    "dataset.train_img_dir": f"{S3_TRAIN}/images.tar.gz",
    "dataset.train_ann_path": f"{S3_TRAIN}/annotations.json",
    "dataset.val_img_dir": f"{S3_EVAL}/images.tar.gz",
    "dataset.val_ann_path": f"{S3_EVAL}/annotations.json",
}

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

{
    "evaluate.checkpoint": "<selected train/AutoML checkpoint>",
    "dataset.val_img_dir": f"{S3_EVAL}/images.tar.gz",
    "dataset.val_ann_path": f"{S3_EVAL}/annotations.json",
}

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

{
    "inference.checkpoint": "<selected train/AutoML checkpoint>",
    "inference.img_dir": f"{S3_EVAL}/images.tar.gz",
    "inference.ann_path": f"{S3_EVAL}/annotations.json",
}

チェックポイントに依存するアクションでは、 references/skill_info.yaml で宣言されたモデルリゾルバーを使用してください。ユーザーが要求した正確な エポック/ステップのチェックポイントを選択するか、ベストチェックポイントのアクションが要求された場合は ベストチェックポイントを選択してください。 mal_model_latest.pth シンボリックリンクは、ユーザーが明示的に 最新のチェックポイントを要求した場合にのみ適切です。

Eval データセット

任意。train パスとともに設定される val 画像とアノテーション。

重要なパラメータ

  • model.arch: ViT-MAE バックボーンのバリアント。デフォルトは vit-mae-base/16。 vit-deit-tiny/16 は避けてください。現在のランタイムは tiny ViT バリアントを拒否します。
  • train.lr: 学習率。デフォルトは 1e-6 (非常に低い — ViT のファインチューニング用)。
  • dataset.crop_size: 学習時のクロップサイズ。デフォルトは 512。model.crop_size ではなく このキーを使用してください。
  • train.warmup_epochs: 完全な学習率に到達する前のウォームアップエポック数。
  • dataset.load_mask: アノテーションに事前計算されたセグメンテーション マスクが含まれているかどうか。ボックスを含むが segmentation フィールドを含まない COCO アノテーションで、 マスクのグラウンドトゥルースなしで学習または推論を実行する場合は、これを false に設定してください。すべてのアノテーションにセグメンテーション データが含まれている場合は true のままにしてください。mIoU で選択する AutoML 検証には、セグメンテーションのグラウンド トゥルースと dataset.load_mask: true が必要です。bbox のみの評価では非有限の mIoU が出力されるため、有効な AutoML 目的関数として受け入れてはなりません。

AutoML / HPO に関する注意

MAL AutoML の起動では、デフォルトのスモーク探索空間を狭く保ち、 automl_hyperparameters=["train.lr", "train.wd"] を渡してください。ViT-MAE ファインチューニングの デフォルトを中心に、保守的なベイズ範囲を使用してください。たとえば train.lr は 1e-7 から 1e-5、train.wd は 1e-5 から 1e-2 です。 パッケージ化された train スキーマは、これら 2 つのパラメータをデフォルトの AutoML パラメータとしてマークしています。バンドルされた config モジュールから MAL 検索メタデータを依然として導出するランタイムを使用する場合は、それらを明示的に渡してください。

マルチ GPU / マルチノード

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

Spec キー説明デフォルト
train.num_gpusGPU 数1
train.gpu_idsGPU デバイスインデックス[0]
train.num_nodesノード数1
  • マルチ GPU 戦略: ddp_find_unused_parameters_true
  • fsdp サポートなし
  • LR 自動スケーリング: lr = lr * num_devices * batch_size (学習率はデバイス数とバッチサイズによって自動的にスケーリングされます)

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

ハードウェア

最低 1 GPU、推奨 2 GPU。GPU あたり 24GB 以上 (A100 推奨) の VRAM。crop_size=512 の ViT-MAE バックボーンには 24GB 以上の GPU メモリが必要です。

エラーパターン

CUDA out of memory:dataset.crop_size を小さくする(512 -> 384 -> 256)か、より小さい ViT-MAE バリアント(base と large のうち base)を使用してください。

MALModelConfig にキー crop_size がない:クロップサイズの上書き設定が model.crop_size の下に配置されています。dataset.crop_size に移動してください。

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

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

TAO Core mal.config.json の推論マッピング:

アクションSpec フィールド推論関数意味
評価evaluate.checkpointparent_model親ジョブの結果フォルダーから推論するモデルファイル
評価results_diroutput_dir現在のジョブの結果ディレクトリ
推論inference.checkpointparent_model親ジョブの結果フォルダーから推論するモデルファイル
推論inference.label_dump_pathcreate_inference_result_file_malMAL 推論 JSON パス
推論results_diroutput_dir現在のジョブの結果ディレクトリ
トレーニングtrain.pretrained_model_pathptm_if_no_resume_model再開しない場合に使用する任意の事前学習済みモデル
トレーニングtrain.resume_training_checkpoint_pathresume_model再開実行用の正確なチェックポイント
トレーニングresults_diroutput_dir現在のジョブの結果ディレクトリ

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

コメント