tao-train-oneformer

OneFormer

ユニバーサル画像セグメンテーション用の OneFormer です。タスク条件付きクエリを使用する単一のアーキテクチャで、パノプティック、インスタンス、セマンティックセグメンテーションを統合します。

train.pretrained_backbone および／または train.pretrained_model を設定してください。

TAO Deploy の TensorRT アクション（gen_trt_engine、TensorRT の evaluate、TensorRT の inference）については、最初に references/tao-deploy-oneformer.md をお読みください。デプロイ仕様テンプレートは、このスキルの 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 経由で実行してください。データセット、仕様、出力ディレクトリ、GPU／プラットフォーム設定、親チェックポイント、automl_policy に関するワークフロー／アプリケーションのオーバーライドは保持してください。automl_policy: off の場合、またはパッケージ化された train スキーマ／テンプレートがない場合に限り、モデルの直接トレーニングを使用してください。スキーマがない場合は、AutoML は有効ですが、スキーマが生成されるまではこのモデルで実行できないことを報告してください。

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

トレーニング要件

• データセットタイプ： セグメンテーション
• 形式： coco_panoptic、coco
• モニタリング指標： mIoU

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

アクション	Spec キー	ソース	ファイル	リスト？
評価	dataset.train.images	train_datasets	images.tar.gz	いいえ
評価	dataset.label_map	train_datasets	label_map.json	いいえ
評価	dataset.train.annotations	train_datasets	annotations.json	いいえ
評価	dataset.train.panoptic	train_datasets	images_panoptic.tar.gz	いいえ
評価	dataset.val.images	eval_dataset	images.tar.gz	いいえ
評価	dataset.val.annotations	eval_dataset	annotations.json	いいえ
評価	dataset.val.panoptic	eval_dataset	images_panoptic.tar.gz	いいえ
評価	dataset.test.images	eval_dataset	images.tar.gz	いいえ
評価	dataset.test.annotations	eval_dataset	annotations.json	いいえ
評価	dataset.test.panoptic	eval_dataset	images_panoptic.tar.gz	いいえ
推論	dataset.train.images	train_datasets	images.tar.gz	いいえ
推論	dataset.label_map	train_datasets	coco_panoptic: label_map_panoptic.json; *: label_map.json	いいえ
推論	dataset.train.annotations	train_datasets	annotations.json	いいえ
推論	dataset.train.panoptic	train_datasets	images_panoptic.tar.gz	いいえ
推論	dataset.val.images	eval_dataset	images.tar.gz	いいえ
推論	dataset.val.annotations	eval_dataset	annotations.json	いいえ
推論	dataset.val.panoptic	eval_dataset	images_panoptic.tar.gz	いいえ
推論	dataset.test.images	eval_dataset	images.tar.gz	いいえ
量子化	dataset.train.images	train_datasets	images.tar.gz	いいえ
量子化	dataset.train.annotations	train_datasets	annotations.json	いいえ
量子化	dataset.label_map	train_datasets	label_map.json	いいえ
量子化	dataset.train.panoptic	train_datasets	images_panoptic.tar.gz	いいえ
量子化	dataset.val.images	eval_dataset	images.tar.gz	いいえ
量子化	dataset.val.annotations	eval_dataset	annotations.json	いいえ
量子化	dataset.val.panoptic	eval_dataset	images_panoptic.tar.gz	いいえ
量子化	dataset.test.images	eval_dataset	images.tar.gz	いいえ
量子化	dataset.quant_calibration_dataset.images_dir	train_datasets	images.tar.gz	いいえ
トレーニング	dataset.train.images	train_datasets	images.tar.gz	いいえ
トレーニング	dataset.train.annotations	train_datasets	annotations.json	いいえ
トレーニング	dataset.label_map	train_datasets	label_map.json	いいえ
トレーニング	dataset.train.panoptic	train_datasets	images_panoptic.tar.gz	いいえ
トレーニング	dataset.val.images	eval_dataset	images.tar.gz	いいえ
トレーニング	dataset.val.annotations	eval_dataset	annotations.json	いいえ
トレーニング	dataset.val.panoptic	eval_dataset	images_panoptic.tar.gz	いいえ
トレーニング	dataset.test.images	eval_dataset	images.tar.gz	いいえ

一般的な Spec の上書き

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

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

トレーニング（必須データソース）：

{
    "train.num_gpus": 1,
    "train.num_epochs": 10,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "model.sem_seg_head.num_classes": 133,
    "dataset.contiguous_id": True,
    "train.precision": "32",
    "dataset.train.images": f"{S3_TRAIN}/images.tar.gz",
    "dataset.train.annotations": f"{S3_TRAIN}/annotations.json",
    "dataset.label_map": f"{S3_TRAIN}/label_map.json",
    "dataset.train.panoptic": f"{S3_TRAIN}/images_panoptic.tar.gz",
    "dataset.val.images": f"{S3_EVAL}/images.tar.gz",
    "dataset.val.annotations": f"{S3_EVAL}/annotations.json",
    "dataset.val.panoptic": f"{S3_EVAL}/images_panoptic.tar.gz",
    "dataset.test.images": f"{S3_EVAL}/images.tar.gz",
}

評価（必須データソース）：

{
    "model.sem_seg_head.num_classes": 133,
    "dataset.contiguous_id": True,
    "dataset.train.images": f"{S3_TRAIN}/images.tar.gz",
    "dataset.label_map": f"{S3_TRAIN}/label_map.json",
    "dataset.train.annotations": f"{S3_TRAIN}/annotations.json",
    "dataset.train.panoptic": f"{S3_TRAIN}/images_panoptic.tar.gz",
    "dataset.val.images": f"{S3_EVAL}/images.tar.gz",
    "dataset.val.annotations": f"{S3_EVAL}/annotations.json",
    "dataset.val.panoptic": f"{S3_EVAL}/images_panoptic.tar.gz",
    "dataset.test.images": f"{S3_EVAL}/images.tar.gz",
    "dataset.test.annotations": f"{S3_EVAL}/annotations.json",
    "dataset.test.panoptic": f"{S3_EVAL}/images_panoptic.tar.gz",
}

エクスポート：

{
    "model.sem_seg_head.num_classes": 133,
    "model.export": True,
}

推論（必須データソース）：

{
    "dataset.train.images": f"{S3_TRAIN}/images.tar.gz",
    "dataset.label_map": {"coco_panoptic": f"{S3_TRAIN}/label_map_panoptic.json; *: label_map.json"},
    "dataset.train.annotations": f"{S3_TRAIN}/annotations.json",
    "dataset.train.panoptic": f"{S3_TRAIN}/images_panoptic.tar.gz",
    "dataset.val.images": f"{S3_EVAL}/images.tar.gz",
    "dataset.val.annotations": f"{S3_EVAL}/annotations.json",
    "dataset.val.panoptic": f"{S3_EVAL}/images_panoptic.tar.gz",
    "dataset.test.images": f"{S3_EVAL}/images.tar.gz",
}

量子化（必須データソース）：

{
    "dataset.train.images": f"{S3_TRAIN}/images.tar.gz",
    "dataset.train.annotations": f"{S3_TRAIN}/annotations.json",
    "dataset.label_map": f"{S3_TRAIN}/label_map.json",
    "dataset.train.panoptic": f"{S3_TRAIN}/images_panoptic.tar.gz",
    "dataset.val.images": f"{S3_EVAL}/images.tar.gz",
    "dataset.val.annotations": f"{S3_EVAL}/annotations.json",
    "dataset.val.panoptic": f"{S3_EVAL}/images_panoptic.tar.gz",
    "dataset.test.images": f"{S3_EVAL}/images.tar.gz",
    "dataset.quant_calibration_dataset.images_dir": f"{S3_TRAIN}/images.tar.gz",
}

評価データセット

任意です。検証データは、データセット設定でトレーニングデータと合わせて設定します。

重要なパラメーター

• model.sem_seg_head.num_classes：セグメンテーションクラスの数です。デフォルトは133（COCO panoptic）です。
• model.backbone.name：デフォルトはD2SwinTransformer（Swinベース）です。デフォルトではembed_dim=192、depths=[2,2,18,2]です。
• train.num_epochs：デフォルトは50です。ほとんどのTAOモデルより大幅に多くなっています。OneFormerの収束には、より多くのエポックが必要です。
• train.optim.lr：学習率です。デフォルトは1e-5です。Mask2Formerの2e-4より低い値です。
• model.task_toggling：特定のタスク（semantic_on、instance_on、panoptic_on）を有効または無効にします。
• export.task：エクスポートするタスクのモードです。選択肢はsemantic、instance、panopticです。デフォルトはsemanticです。エクスポート入力のデフォルトは640x640です。
• inference.mode：推論モードです。選択肢はsemantic、instance、panopticです。デフォルトはsemanticです。image_sizeのデフォルトは[1024, 1024]です。
• evaluate.iou_per_class：評価時にクラスごとのIoUをレポートします。デフォルトはTrueです。

マルチGPU／マルチノード

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

Spec Key	説明	デフォルト
train.num_gpus	GPUの数	1
train.gpu_ids	GPUデバイスのインデックス	[0]
train.num_nodes	ノードの数	1

• find_unused_parameters=True、gradient_as_bucket_view=True、process_group_backend="nccl"を指定したDDPStrategyを使用します
• sync_batchnormは常に有効です
• FSDPはサポートされていません。DDPのみです

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

ハードウェア

GPUは最低2基、推奨4基です。GPUごとに24GB以上（A100を推奨）のVRAMが必要です。OneFormerはMask2Formerと同様、メモリーを大量に消費します。デフォルトのbatch_sizeは1です。妥当な速度で学習するには、特に50エポックの場合、マルチGPUが必要です。

エラーパターン

CUDAのメモリー不足：batch_sizeはすでに1です。画像解像度を下げるか、より小さいSwin設定を使用してください。

学習が遅い：batch_size=1でデフォルトの50エポックを実行すると、シングルGPUでは時間がかかります。マルチGPUの分散学習を使用してください。

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

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

TAO Core の oneformer.config.json からの推論マッピング：

アクション	Spec フィールド	推論関数	意味
評価	encryption_key	key	暗号化キー
評価	evaluate.checkpoint	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
評価	evaluate.trt_engine	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
評価	results_dir	output_dir	現在のジョブの結果ディレクトリ
エクスポート	encryption_key	key	暗号化キー
エクスポート	export.checkpoint	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
エクスポート	export.onnx_file	create_onnx_file	出力 ONNX パス
エクスポート	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	現在のジョブの結果ディレクトリ
推論	encryption_key	key	暗号化キー
推論	inference.checkpoint	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
推論	inference.trt_engine	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
推論	results_dir	output_dir	現在のジョブの結果ディレクトリ
量子化	encryption_key	key	暗号化キー
量子化	quantize.model_path	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
量子化	results_dir	output_dir	現在のジョブの結果ディレクトリ
学習	encryption_key	key	暗号化キー
学習	results_dir	output_dir	現在のジョブの結果ディレクトリ
学習	train.pretrained_backbone	{'link': 'https://github.com/SwinTransformer/storage/releases/download/v1.0.8/swin_tiny_patch4_window7_224_22k.pth', 'destination_path': '/ptm/mask2former/swin_tiny_patch4_window7_224_22k/swin_tiny_patch4_window7_224_22k.pth'}	{'link': 'https://github.com/SwinTransformer/storage/releases/download/v1.0.8/swin_tiny_patch4_window7_224_22k.pth', 'destination_path': '/ptm/mask2former/swin_tiny_patch4_window7_224_22k/swin_tiny_patch4_window7_224_22k.pth'}
学習	train.pretrained_model	ptm_if_no_resume_model	再開用チェックポイントがない場合の PTM
学習	train.resume_training_checkpoint_path	resume_model	現在のジョブの結果フォルダーから推論されるモデルファイル

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

デプロイ

• tao-deploy-oneformer — TAO Deploy を使用した、TensorRT エンジンの生成、TensorRT による評価、および TensorRT による推論の OneFormer デプロイワークフロー。