tao-train-segformer について
SegFormer によるセマンティックセグメンテーション向けの TAO スキルです。AutoML 対応の学習に加え、推論、エクスポート、量子化を扱い、TensorRT 操作には別のデプロイフローを使います。
主な機能
- 01SegFormer によるセマンティックセグメンテーション
- 02AutoML 対応の学習
- 03評価と推論
- 04モデルのエクスポートと量子化
- 05独立した TensorRT デプロイフロー
活用シーン
- 01画像とマスクを使ったセグメンテーションモデルの学習
- 02test_miou によるチェックポイントの評価
- 03セグメンテーションデータセットでの推論
- 04学習済みモデルのエクスポートまたは量子化
インストール
コマンドをコピーしてターミナルで実行します。
npx skills add https:/ / github.com/ nvidia/ skills --skill tao-train-segformerインストール後、Claude Code で/tao-train-segformerを入力して読み込まれたか確認します。
同梱ファイル
SKILL.md と一緒にインストールされる参考資料やスクリプト。
BENCHMARK.mdskill-card.mdskill.oms.sigconfig/skillspector-baseline.yamlevals/evals.jsonreferences/skill_info.yamlspec_template_deploy_evaluate.yamlspec_template_deploy_gen_trt_engine.yamlspec_template_deploy_inference.yamlspec_template_evaluate.yamlspec_template_export.yamlspec_template_inference.yamlspec_template_quantize.yamlspec_template_train.yamltao-deploy-segformer.mdtao-deploy-segformer.skill_info.yamlschemas/evaluate.schema.jsonexport.schema.jsoninference.schema.jsonmanifest.jsonquantize.schema.jsontrain.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_dir | eval_dataset | images/<split> と masks/<split> を含む展開済みルートディレクトリ | いいえ |
| エクスポート | dataset.segment.root_dir | train_datasets | images/<split> と masks/<split> を含む展開済みルートディレクトリ | いいえ |
| 推論 | dataset.segment.root_dir | inference_dataset | images/<split> と masks/<split> を含む展開済みルートディレクトリ | いいえ |
| 量子化 | dataset.segment.root_dir | train_datasets | images/<split> と masks/<split> を含む展開済みルートディレクトリ | いいえ |
| 量子化 | dataset.segment.quant_calibration_dataset.images_dir | calibration_dataset | 展開済み画像ディレクトリ | いいえ |
| トレーニング | dataset.segment.root_dir | train_datasets | images/<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_gpus | GPU 数 | 1 |
train.gpu_ids | GPU デバイスのインデックス | [0] |
train.num_nodes | ノード数 | 1 |
train.sync_batchnorm | GPU 間で 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 からの推論マッピング:
| アクション | 仕様フィールド | 推論関数 | 意味 |
|---|---|---|---|
| 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 | 現在のジョブの結果ディレクトリ |
| quantize | encryption_key | key | 暗号化キー |
| quantize | quantize.model_path | parent_model | 親ジョブの結果フォルダーから推論されるモデルファイル |
| quantize | results_dir | output_dir | 現在のジョブの結果ディレクトリ |
| train | encryption_key | key | 暗号化キー |
| train | model.backbone.pretrained_backbone_path | ptm_if_no_resume_model | 再開用チェックポイントがない場合の PTM |
| 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 に戻したり、チェックポイントパスを推測するように生成済みのランナースクリプトを修正したりしないでください。