tao-train-rtdetr

RT-DETR

RT-DETR（Real-Time DEtection TRansformer）は、2D 物体検出向けです。高い精度を維持しながら、リアルタイム推論を実現するように設計されています。デプロイの最適化に向けて、蒸留と量子化をサポートしています。

バックボーンの重みには model.pretrained_backbone_path を、モデル全体の重みには train.pretrained_model_path を設定してください。

TAO Deploy の TensorRT アクション（gen_trt_engine、TensorRT の evaluate、TensorRT の inference）を実行する場合は、先に references/tao-deploy-rtdetr.md をお読みください。デプロイ仕様テンプレートは、このスキルの 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/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 の上書きによって、モデルのメタデータが変更されることはありません。

トレーニング要件

• データセットタイプ： object_detection
• 形式： coco, coco_raw
• 監視指標： val_mAP50

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

アクション	仕様キー	ソース	ファイル	リスト？
distill	dataset.train_data_sources	train_datasets	image_dir: images.tar.gz, json_file: annotations.json	はい
distill	dataset.val_data_sources	eval_dataset	image_dir: images.tar.gz, json_file: annotations.json	いいえ
evaluate	dataset.test_data_sources	eval_dataset	image_dir: images.tar.gz, json_file: annotations.json	いいえ
gen_trt_engine	gen_trt_engine.tensorrt.calibration.cal_image_dir	calibration_dataset	images.tar.gz	はい
inference	dataset.infer_data_sources	inference_dataset	image_dir: images.tar.gz, classmap: label_map.txt	いいえ
quantize	dataset.train_data_sources	train_datasets	image_dir: images.tar.gz, json_file: annotations.json	はい
quantize	dataset.val_data_sources	eval_dataset	image_dir: images.tar.gz, json_file: annotations.json	いいえ
quantize	dataset.quant_calibration_data_sources	train_datasets	image_dir: images.tar.gz, json_file: annotations.json	いいえ
train	dataset.train_data_sources	train_datasets	image_dir: images.tar.gz, json_file: annotations.json	はい
train	dataset.val_data_sources	eval_dataset	image_dir: images.tar.gz, json_file: annotations.json	いいえ

一般的な仕様のオーバーライド

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

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

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

{
    "train.num_epochs": 10,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "train.num_gpus": 1,
    "dataset.num_classes": "<num_classes> + 1",
    "dataset.train_data_sources": [{"image_dir": f"{S3_TRAIN}/images.tar.gz", "json_file": f"{S3_TRAIN}/annotations.json"}],
    "dataset.val_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "json_file": f"{S3_EVAL}/annotations.json"},
}

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

{
    "dataset.num_classes": "<num_classes> + 1",
    "dataset.test_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "json_file": f"{S3_EVAL}/annotations.json"},
}

export：

{
    "dataset.num_classes": "<num_classes> + 1",
    "export.input_height": 640,
    "export.input_width": 640,
}

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

{
    "dataset.num_classes": "<num_classes> + 1",
    "quantize.layers": [
        {
            "module_name": "*",
            "weights": {
                "dtype": "float8_e4m3fn"
            },
            "activations": {
                "dtype": "float8_e4m3fn"
            }
        }
    ],
    "dataset.train_data_sources": [{"image_dir": f"{S3_TRAIN}/images.tar.gz", "json_file": f"{S3_TRAIN}/annotations.json"}],
    "dataset.val_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "json_file": f"{S3_EVAL}/annotations.json"},
    "dataset.quant_calibration_data_sources": {"image_dir": f"{S3_TRAIN}/images.tar.gz", "json_file": f"{S3_TRAIN}/annotations.json"},
}

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

{
    "gen_trt_engine.tensorrt.data_type": "FP16",
    "gen_trt_engine.tensorrt.calibration.cal_image_dir": [f"{S3_TRAIN}/images.tar.gz"],
}

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

{
    "dataset.num_classes": "<num_classes> + 1",
    "dataset.infer_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "classmap": f"{S3_EVAL}/label_map.txt"},
}

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

{
    "dataset.train_data_sources": [{"image_dir": f"{S3_TRAIN}/images.tar.gz", "json_file": f"{S3_TRAIN}/annotations.json"}],
    "dataset.val_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "json_file": f"{S3_EVAL}/annotations.json"},
}

評価データセット

任意です。指定すると、各チェックポイントで検証 mAP が提供されます。

重要なパラメーター

• dataset.num_classes：クラス数です。デフォルトは80（MSCOCOの80クラス）です。データセットのアノテーションと一致させる必要があります。
• model.backbone：デフォルトはresnet_50です。サポート対象：ResNetの各種バリアント、ConvNeXt、FAN、EfficientViT。RT-DETRは、より軽量なバックボーンでリアルタイム処理に最適化されています。
• train.optim.lr：学習率です。デフォルトは1e-4（DINOの2e-4より低い値）です。lr_backboneのデフォルトは1e-5です。
• dataset.augmentation.train_spatial_size：学習時の入力サイズです。デフォルトは[640, 640]です。DINOのマルチスケール（最大1333）より小さい値です。RT-DETRの速度を高める重要な要素です。
• model.num_feature_levels：デフォルトは3（DINOは4）です。return_interm_indicesは[1,2,3]です。
• train.enable_ema：指数移動平均です。デフォルトはFalseです。有効にすると、収束がより安定する可能性があります。
• dataset.remap_mscoco_category：デフォルトはFalseです。91から80へのカテゴリ ID マッピングを行う、元のMSCOCOデータセットの場合にのみTrueに設定してください。

マルチGPU／マルチノード

起動方法： torchrun（LIGHTNING_EXCLUDED_NETWORK）。エントリーポイントはtorchrun --nnodes=N --nproc-per-node=M train.pyを実行します。通常のpythonではありません。

仕様キー	説明	デフォルト
train.num_gpus	ノードあたりのGPU数	1
train.gpu_ids	GPUデバイスのインデックス	[0]
train.num_nodes	ノード数	1
train.distributed_strategy	ddpまたはfsdp	ddp

• CUDA_VISIBLE_DEVICESが明示的に設定されます（Lightningが管理するモデルではTAO_VISIBLE_DEVICESが使用されます）
• アクティベーションチェックポイントを使用するddp：find_unused_parameters=False
• 使用しない場合のddp：find_unused_parameters=True
• fsdpをサポートしており、FP16が強制されます

マルチノード環境変数（オーケストレーターによって設定）：

変数	目的
WORLD_SIZE	ノード数（マルチノードモードを有効化）
NODE_RANK	このノードのランク（0から始まるインデックス）
MASTER_ADDR	ランク0ノードのIP
MASTER_PORT	ランク0のポート（デフォルトは29500）
NUM_GPU_PER_NODE	ノードあたりのGPU数（デフォルト：表示されているすべてのGPU）

重要： RANKが未設定の場合、NODE_RANKがRANKにコピーされます。これはtorchrunのマルチノード実行に必要です。

エクスポート／TRT のデフォルト設定

• エクスポート入力：640x640、opset 17
• TRT のデータ型：FP32、FP16、INT8
• TRT のワークスペース：1024 MB
• TRT の max_batch_size：4

TAO Deploy の全リファレンス：tao-deploy-rtdetr。

蒸留

RT-DETR は、教師モデルを使用した知識蒸留に対応しています。teacher model のパスと蒸留バインディングの設定を指定して distill アクションを実行する必要があります。

ハードウェア

GPU は最低 1 基、推奨 2 基です。各 GPU に 16GB 以上（V100 または A100）の VRAM が必要です。RT-DETR は入力サイズが小さく（640x640）、特徴レベルも少ないため、DINO／GDINO よりメモリ効率に優れています。小規模から中規模のデータセットであれば、単一 GPU でも十分に学習できます。

エラーのパターン

CUDA のメモリ不足：batch_size を減らしてください。640x640 の RT-DETR は 1333px の DINO より軽量ですが、16GB の GPU では batch_size が 8 を超えるとメモリ不足になる場合があります。

num_classes の不一致：RT-DETR のデフォルトは 80（DINO のように 91 ではありません）です。dataset.num_classes がアノテーションのカテゴリ数と一致していることを確認してください。

return_interm_indices と num_feature_levels の不一致：デフォルトは [1,2,3]、num_feature_levels=3 です。変更する場合は、両者の整合性を保つ必要があります。

Spec Param／親モデルの推論

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

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

アクション	Spec フィールド	推論関数	意味
distill	distill.pretrained_teacher_model_path	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
distill	encryption_key	key	暗号化キー
distill	results_dir	output_dir	現在のジョブの結果ディレクトリ
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.tensorrt.calibration.cal_cache_file	create_cal_cache	キャリブレーションキャッシュのパス
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.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 に戻したり、チェックポイントのパスを推測するように生成済みランナースクリプトを修正したりしないでください。