mcpField
T
tao-train-segformer

tao-train-segformer

github.com/nvidia開発3.8k1.6k 回インストール今日更新
リポジトリ

tao-train-segformer について

SegFormer によるセマンティックセグメンテーション向けの TAO スキルです。AutoML 対応の学習に加え、推論、エクスポート、量子化を扱い、TensorRT 操作には別のデプロイフローを使います。

主な機能

  1. 01SegFormer によるセマンティックセグメンテーション
  2. 02AutoML 対応の学習
  3. 03評価と推論
  4. 04モデルのエクスポートと量子化
  5. 05独立した TensorRT デプロイフロー

活用シーン

  1. 01画像とマスクを使ったセグメンテーションモデルの学習
  2. 02test_miou によるチェックポイントの評価
  3. 03セグメンテーションデータセットでの推論
  4. 04学習済みモデルのエクスポートまたは量子化

インストール

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

$npx skills add https://github.com/nvidia/skills --skill tao-train-segformer

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

同梱ファイル

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

BENCHMARK.md
skill-card.md
skill.oms.sig
config/
skillspector-baseline.yaml
evals/
evals.json
references/
skill_info.yaml
spec_template_deploy_evaluate.yaml
spec_template_deploy_gen_trt_engine.yaml
spec_template_deploy_inference.yaml
spec_template_evaluate.yaml
spec_template_export.yaml
spec_template_inference.yaml
spec_template_quantize.yaml
spec_template_train.yaml
tao-deploy-segformer.md
tao-deploy-segformer.skill_info.yaml
schemas/
evaluate.schema.json
export.schema.json
inference.schema.json
manifest.json
quantize.schema.json
train.schema.json

詳細

SegFormer

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

セマンティックセグメンテーション用の SegFormer です。階層的な特徴抽出を備えた軽量なトランスフォーマーベースのアーキテクチャです。リアルタイムのセグメンテーションタスクに適しています。

バックボーンの重みには model.backbone.pretrained_backbone_path を設定してください。

TAO Deploy の TensorRT アクション(gen_trt_engine、TensorRT の evaluate、TensorRT の inference)については、まず references/tao-deploy-segformer.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 が利用できることを前提にしないでください。メンテナーはスキルバンクをパッケージ化する前に、スキーマとテンプレートを再生成してください。

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、deploy フローなどの train 以外のアクションは、このモデルスキル内で処理します。実行ごとの automl_policy オーバーライドによって、モデルのメタデータが変わることはありません。

サポートされるアクション

パッケージ化された SegFormer PyT CLI は、train、evaluate、export、inference、quantize、default_specs をサポートします。このモデルスキルが公開するのは train、evaluate、export、inference、quantize です。再開/再学習は、train.resume_training_checkpoint_path を指定して train を実行します。

親 PyT CLI は gen_trt_engine を公開していません。TensorRT エンジンの生成、TensorRT 評価、TensorRT 推論には models/segformer/deploy を使用してください。

学習要件

  • データセットの種類: セグメンテーション
  • 形式: unet
  • AutoML の学習メトリック: val_miou(最大化)。
  • 単体評価メトリック: test_miou。AutoML のランキングには学習ステータスの val_miou を使用し、評価器の test_miou はチェックポイントの検証にのみ使用してください。
アクション別データセット要件
アクションSpec キーソースファイルリスト形式?
評価dataset.segment.root_direval_datasetimages/<split> と masks/<split> を含む展開済みルートディレクトリいいえ
エクスポートdataset.segment.root_dirtrain_datasetsimages/<split> と masks/<split> を含む展開済みルートディレクトリいいえ
推論dataset.segment.root_dirinference_datasetimages/<split> と masks/<split> を含む展開済みルートディレクトリいいえ
量子化dataset.segment.root_dirtrain_datasetsimages/<split> と masks/<split> を含む展開済みルートディレクトリいいえ
量子化dataset.segment.quant_calibration_dataset.images_dircalibration_dataset展開済み画像ディレクトリいいえ
トレーニングdataset.segment.root_dirtrain_datasetsimages/<split> と masks/<split> を含む展開済みルートディレクトリいいえ
一般的な Spec のオーバーライド

データソースのオーバーライドはすべてのアクションで必須です。エージェントは、上記のアクション別データセット要件の表に基づいてデータソースのパスを構築し、spec_overrides に含める必要があります。

SEG_TRAIN_ROOT = "/data/segformer/train"
SEG_EVAL_ROOT = "/data/segformer/eval"
SEG_INFER_ROOT = "/data/segformer/infer"
CAL_IMAGES = f"{SEG_TRAIN_ROOT}/images/train"

トレーニング(必須データソース):

{
    "train.num_gpus": 1,
    "train.num_epochs": 10,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "dataset.segment.batch_size": 4,
    "dataset.segment.root_dir": SEG_TRAIN_ROOT,
}

評価(必須データソース):

{
    "evaluate.batch_size": 4,
    "dataset.segment.root_dir": SEG_EVAL_ROOT,
    "evaluate.checkpoint": CHECKPOINT,
}

推論(必須データソース):

{
    "dataset.segment.batch_size": 1,
    "dataset.segment.root_dir": SEG_INFER_ROOT,
    "inference.checkpoint": CHECKPOINT,
}

エクスポート(必須データソース):

{
    "dataset.segment.root_dir": SEG_TRAIN_ROOT,
    "export.checkpoint": CHECKPOINT,
    "export.input_height": 256,
    "export.input_width": 256,
    "export.onnx_file": ONNX_FILE,
}

量子化(必須データソース):

{
    "dataset.segment.root_dir": SEG_TRAIN_ROOT,
    "dataset.segment.quant_calibration_dataset.images_dir": CAL_IMAGES,
    "quantize.model_path": CHECKPOINT,
}

ソースデータセットが images/*.tar.gz と masks/*.tar.gz に分かれたアーカイブとして提供される場合は、起動前に展開し、root_dir に images/train、images/val、images/test、masks/train、masks/val などのディレクトリが含まれるようにしてください。tarball だけが格納されたままのアーカイブ展開用フォルダーを dataset.segment.root_dir に指定しないでください。

評価データセット

任意です。検証データは通常、root_dir の構造に含まれます。

重要なパラメーター

  • dataset.segment.num_classes:セグメンテーションのクラス数です。デフォルトは 2(バイナリ)です。マスクのアノテーションに含まれるクラス数と一致させる必要があります。
  • model.backbone.type:デフォルトは fan_small_12_p4_hybrid です。FAN の各種バリエーション、SegFormer MIT の各種バリエーションなどに対応しています。
  • dataset.segment.root_dir:セグメンテーションデータセットのルートディレクトリです。
  • dataset.segment.img_size:入力画像のサイズです。デフォルトは 256 です。メモリ使用量は増加しますが、より細かいセグメンテーションを行う場合は大きくしてください。
  • train.optim.lr:学習率です。デフォルトは 6e-5 です。
  • model.freeze_backbone:トレーニング中にバックボーンを凍結するかどうかを指定します。データが限られている場合のファインチューニングに役立ちます。
  • dataset.segment.batch_size:GPU ごとのバッチサイズです。デフォルトは 8 です。
  • dataset.segment.label_transform:ラベル変換が不要な場合は、文字列 "None" を使用してください。JSON/YAML の null に設定しないでください。厳密なスキーママージでは、このフィールドは文字列の列挙値として扱われます。
  • dataset.segment.palette:グレースケールマスクの場合は、RGB の各エントリに整数を 1 つ指定します。たとえば rgb: [85] のようにします。ユーザーが変換を明示的に要求しない限り、データセットの実際のラベル ID とクラス名を維持し、正規化しないでください。

マルチ GPU/マルチノード

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

Spec キー説明デフォルト
train.num_gpusGPU 数1
train.gpu_idsGPU デバイスのインデックス[0]
train.num_nodesノード数1
train.sync_batchnormGPU 間で BN を同期するかどうか設定可能
train.use_distributed_sampler分散サンプラーを使用するかどうか設定可能
  • マルチ GPU 戦略:ddp_find_unused_parameters_true
  • fsdp はサポートされていません

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

ハードウェア

最小 1 GPU、推奨 2 GPU です。各 GPU に 16GB 以上(V100 または A100)の VRAM が必要です。SegFormer は比較的軽量です。デフォルトの img_size=256 はメモリにやさしい設定です。メモリ使用量と速度を犠牲にして、img_size を大きくすると高解像度になります。

エラーパターン

CUDA のメモリ不足:batch_size または img_size を小さくしてください。SegFormer のメモリ使用量は画像サイズの二乗に比例します。

num_classes の不一致:dataset.segment.num_classes が、マスクのアノテーションに含まれる実際のクラス数と一致していることを確認してください。

セグメンテーションの学習では TensorBoard はサポートされません:train.tensorboard.enabled: false のままにしてください。SegFormer の学習エントリーポイントでは、セグメンテーションでの TensorBoard 可視化はサポートされていないとアサートされます。そのため、AutoML のメトリクスを抽出する目的で TensorBoard を有効にしないでください。代わりに、ログ解析または学習後の評価器を使用してください。

AutoML のメトリクス抽出:SegFormer の学習ステータスファイルには、val_loss、val_acc、その他の検証 KPI とともに val_miou が報告されます。デフォルトの AutoML 学習起動では、direction: maximize を指定して val_miou を最適化してください。デフォルトのモデル呼び出しでは val_loss を最適化しないでください。

AutoML または長時間のセグメンテーション探索では、まず results_dir/train/status.json から val_miou を読み取ってください。ラッパーが最終的な 失敗を報告した場合でも、構造化されたステータスファイルが設定済みの学習予算に 到達しており、有効な val_miou が含まれているなら、測定値を破棄せず、 ラッパーの失敗を明記したうえで取得できたメトリクスを報告してください。

高解像度のカスタムセグメンテーションターゲットでは、データセットのパスを実行ごとの 入力として指定してください。再利用可能なこのスキルに、顧客やユーザー固有の ルートを追加しないでください。固定されたフル予算での探索をユーザーが求めている場合は、 ブラケットアルゴリズム(asha、bohb、dehb、hyperband、hyperband_es、pbt)が 一部の推奨設定で train.num_epochs を意図的に少なくする場合があることに注意してください。 すべての推奨設定でエポック数をフルに実行する必要があるなら、ベイズ最適化/BFBO を使用するか、 予算を固定してください。

チェックポイントの受け渡し:評価/エクスポート/推論/量子化/再開では、最良の AutoML 子ジョブの results_dir/train/ フォルダーに対してチェックポイントリゾルバーを使用し、アクションに適した model_epoch_*.pth チェックポイント(例:model_epoch_000_step_00010.pth)を選択してください。SegFormer は segformer_model_latest.pth も書き出す場合がありますが、呼び出し元が latest を明示的に要求した場合に限り、 これを使用してください。後続のアクションでも dataset.segment.num_classes、dataset.segment.img_size、 dataset.segment.root_dir のオーバーライドを維持してください。

チェックポイントからの再開/再学習:再開には train.resume_training_checkpoint_path を使用します。 推測した model.pth のパスではなく、前回の学習出力から解決した正確なチェックポイントを 渡してください。1 エポックの再開実行では、新しい結果ディレクトリに次のチェックポイントが 生成されます(例:model_epoch_001_step_00020.pth)。

エクスポート/TensorRT の形状の整合:学習済みモデルとデプロイ仕様が別の解像度で 検証されていない限り、export.input_height と export.input_width を dataset.segment.img_size と揃えてください。パッケージに含まれる 新規インストール用のパスは 256x256 で検証されており、デフォルトの SegFormer データセットおよびデプロイテンプレートに一致します。

親モデルの segformer gen_trt_engine が PyT CLI に拒否される場合:検証済みの 7.0.0 PyT コンテナーでは、segformer gen_trt_engine は有効な親モデルのサブタスクではありません。TensorRT エンジンの生成、TensorRT 評価、TensorRT 推論には、SegFormer デプロイのワークフロー(references/tao-deploy-segformer.md)を使用してください。

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

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

TAO Core の segformer.config.json からの推論マッピング:

アクション仕様フィールド推論関数意味
evaluateencryption_keykey暗号化キー
evaluateevaluate.checkpointparent_model親ジョブの結果フォルダーから推論されるモデルファイル
evaluateevaluate.trt_engineparent_model親ジョブの結果フォルダーから推論されるモデルファイル
evaluateresults_diroutput_dir現在のジョブの結果ディレクトリ
exportencryption_keykey暗号化キー
exportexport.checkpointparent_model親ジョブの結果フォルダーから推論されるモデルファイル
exportexport.onnx_filecreate_onnx_fileONNX 出力パス
exportresults_diroutput_dir現在のジョブの結果ディレクトリ
inferenceencryption_keykey暗号化キー
inferenceinference.checkpointparent_model親ジョブの結果フォルダーから推論されるモデルファイル
inferenceinference.trt_engineparent_model親ジョブの結果フォルダーから推論されるモデルファイル
inferenceresults_diroutput_dir現在のジョブの結果ディレクトリ
quantizeencryption_keykey暗号化キー
quantizequantize.model_pathparent_model親ジョブの結果フォルダーから推論されるモデルファイル
quantizeresults_diroutput_dir現在のジョブの結果ディレクトリ
trainencryption_keykey暗号化キー
trainmodel.backbone.pretrained_backbone_pathptm_if_no_resume_model再開用チェックポイントがない場合の PTM
trainresults_diroutput_dir現在のジョブの結果ディレクトリ
traintrain.pretrained_model_pathptm_if_no_resume_model再開用チェックポイントがない場合の PTM
traintrain.resume_training_checkpoint_pathresume_model現在のジョブの結果フォルダーから推論されるモデルファイル

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

デプロイ

コメント