tao-train-segformer

SegFormer

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

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

TAO DeployのTensorRTアクション（gen_trt_engine、TensorRTのevaluate、TensorRTのinference）については、最初にreferences/tao-deploy-segformer.mdをお読みください。Deploy specテンプレートは、このスキルのreferences/フォルダー内にあり、spec_template_deploy_*.yamlというプレフィックスが付いています。

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

evaluate、inference、export、デプロイフローなどのtrain以外のアクションは、このモデルのスキル内で処理してください。実行ごとのautoml_policyの上書きによって、モデルのメタデータが変わることはありません。

トレーニング要件

• データセットタイプ： segmentation
• 形式： unet
• 監視メトリック： val_miou

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

アクション	Specキー	ソース	ファイル	リスト？
evaluate	dataset.segment.root_dir	eval_dataset		いいえ
export	dataset.segment.root_dir	train_datasets		いいえ
inference	dataset.segment.root_dir	eval_dataset		いいえ
quantize	dataset.segment.root_dir	train_datasets		いいえ
quantize	dataset.segment.quant_calibration_dataset.images_dir	train_datasets		いいえ
train	dataset.segment.root_dir	train_datasets		いいえ

一般的なSpecの上書き

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

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

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": f"{S3_TRAIN}",
}

evaluate（必須データソース）：

{
    "evaluate.batch_size": 4,
    "dataset.segment.root_dir": f"{S3_EVAL}",
}

gen_trt_engine：

{
    "gen_trt_engine.tensorrt.data_type": "fp16",
}

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

{
    "dataset.segment.batch_size": 1,
    "dataset.segment.root_dir": f"{S3_EVAL}",
}

export（必須データソース）：

{
    "dataset.segment.root_dir": f"{S3_TRAIN}",
}

quantize（必須データソース）：

{
    "dataset.segment.root_dir": f"{S3_TRAIN}",
    "dataset.segment.quant_calibration_dataset.images_dir": f"{S3_TRAIN}",
}

評価用データセット

任意です。通常、検証データは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 です。

マルチ GPU／マルチノード

起動方法： Lightning による管理（単一の python プロセスから Lightning がワーカーを起動します）。

Spec Key	説明	デフォルト
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 あたり 16 GB 以上（V100 または A100）の VRAM が必要です。SegFormer は比較的軽量です。デフォルトの img_size=256 はメモリに優しい設定です。解像度を高くするには img_size を大きくしてください。その分、メモリ使用量が増え、速度が低下します。

エラーのパターン

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

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

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

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

TAO Core segformer.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	現在のジョブの結果ディレクトリ
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 に戻したり、チェックポイントのパスを推測するよう生成されたランナースクリプトを修正したりしないでください。

デプロイ

• tao-deploy-segformer — TAO Deploy を使用して TensorRT エンジンを生成し、TensorRT で評価および推論を行う SegFormer のデプロイワークフロー。