tao-train-image-classification

Classification PyT

PyTorch 画像分類です。蒸留と量子化によるデプロイに対応し、幅広いバックボーン（FAN、EfficientNet、ResNet など）をサポートします。

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

TAO Deploy の TensorRT アクション（gen_trt_engine、TensorRT evaluate、TensorRT inference）については、まず references/tao-deploy-image-classification.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 オーバーライドによって、モデルのメタデータは変更されません。

トレーニング要件

• データセットの種類： image_classification
• 形式： classification_pyt
• モニタリング指標： val_acc_1

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

アクション	Spec キー	ソース	ファイル	リスト？
distill	dataset.train_dataset.images_dir	train_datasets	images_train.tar.gz	いいえ
distill	dataset.classes_file	train_datasets	classes.txt	いいえ
distill	dataset.val_dataset.images_dir	eval_dataset	images_val.tar.gz	いいえ
evaluate	dataset.val_dataset.images_dir	eval_dataset	images_val.tar.gz	いいえ
evaluate	dataset.classes_file	eval_dataset	classes.txt	いいえ
evaluate	dataset.test_dataset.images_dir	inference_dataset	images_test.tar.gz	いいえ
export	dataset.root_dir	train_datasets		いいえ
inference	dataset.val_dataset.images_dir	eval_dataset	images_val.tar.gz	いいえ
inference	dataset.classes_file	eval_dataset	classes.txt	いいえ
inference	dataset.test_dataset.images_dir	inference_dataset	images_test.tar.gz	いいえ
quantize	dataset.train_dataset.images_dir	train_datasets	images_train.tar.gz	いいえ
quantize	dataset.classes_file	train_datasets	classes.txt	いいえ
quantize	dataset.val_dataset.images_dir	eval_dataset	images_val.tar.gz	いいえ
quantize	dataset.quant_calibration_dataset.images_dir	calibration_dataset	images_train.tar.gz	いいえ
train	dataset.train_dataset.images_dir	train_datasets	images_train.tar.gz	いいえ
train	dataset.classes_file	train_datasets	classes.txt	いいえ
train	dataset.val_dataset.images_dir	eval_dataset	images_val.tar.gz	いいえ

よく使う Spec オーバーライド

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

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

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

{
    "train.num_epochs": 2,
    "train.validation_interval": 2,
    "train.checkpoint_interval": 2,
    "train.num_gpus": 1,
    "dataset.train_dataset.images_dir": f"{S3_TRAIN}/images_train.tar.gz",
    "dataset.classes_file": f"{S3_TRAIN}/classes.txt",
    "dataset.val_dataset.images_dir": f"{S3_EVAL}/images_val.tar.gz",
}

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

{
    "export.input_height": 224,
    "export.input_width": 224,
    "dataset.root_dir": f"{S3_TRAIN}",
}

gen_trt_engine：

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

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

{
    "dataset.batch_size": 1,
    "dataset.val_dataset.images_dir": f"{S3_EVAL}/images_val.tar.gz",
    "dataset.classes_file": f"{S3_EVAL}/classes.txt",
    "dataset.test_dataset.images_dir": f"{S3_EVAL}/images_test.tar.gz",
}

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

{
    "dataset.train_dataset.images_dir": f"{S3_TRAIN}/images_train.tar.gz",
    "dataset.classes_file": f"{S3_TRAIN}/classes.txt",
    "dataset.val_dataset.images_dir": f"{S3_EVAL}/images_val.tar.gz",
}

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

{
    "dataset.val_dataset.images_dir": f"{S3_EVAL}/images_val.tar.gz",
    "dataset.classes_file": f"{S3_EVAL}/classes.txt",
    "dataset.test_dataset.images_dir": f"{S3_EVAL}/images_test.tar.gz",
}

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

{
    "dataset.train_dataset.images_dir": f"{S3_TRAIN}/images_train.tar.gz",
    "dataset.classes_file": f"{S3_TRAIN}/classes.txt",
    "dataset.val_dataset.images_dir": f"{S3_EVAL}/images_val.tar.gz",
    "dataset.quant_calibration_dataset.images_dir": f"{S3_TRAIN}/images_train.tar.gz",
}

評価データセット

任意です。検証用画像は、トレーニング画像とは別の tar ファイルとして提供されます。

重要なパラメーター

• dataset.num_classes：クラス数です。デフォルトは20です。画像 tar ファイル内のサブディレクトリ数と一致している必要があります。
• model.backbone.type：デフォルトはfan_small_12_p4_hybridです。サポートされているバックボーンとそのhead in_channelsは（model_params_mapping.pyより）次のとおりです。FAN：fan_tiny、fan_small_12_p4_hybrid、fan_base_16_p4_hybrid、fan_large_16_p4_hybrid。GCViT：gcvit_tinyからgcvit_large。FasterViT：fastervit_0からfastervit_6。ViT/EVA/DINO：vit_large_patch14_dinov2、eva02_large_patch14など。SigLIP-CLIPA：ViT-H-14-SigLIP-CLIPA-224など。一部のバックボーンでは、デフォルト以外の入力解像度（384、512、768）が必要です。
• dataset.classes_file：クラス名を列挙したclasses.txtへのパスです。
• train.optim.lr：学習率です。デフォルトは6e-5です。
• dataset.img_size：入力画像サイズです。デフォルトは224です。
• dataset.batch_size：GPUごとのバッチサイズです。デフォルトは8です。

マルチ GPU／マルチノード

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

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

• マルチ 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が必要です。分類処理は一般に軽量です。224x224では、ほとんどのバックボーンが16GB GPU上でbatch_size=8で問題なく動作します。

エラーのパターン

CUDA out of memory：batch_size を減らすか、より小さいバックボーンを使用してください。

num_classes の不一致：dataset.num_classes が、画像の tarball と classes.txt に実際に含まれるクラスディレクトリと一致していることを確認してください。

空のクラスディレクトリ：classes.txt に記載されているすべてのクラスに、対応するサブディレクトリ内の画像が少なくとも 1 枚必要です。

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

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

TAO Core classification_pyt.config.json の推論マッピング：

アクション	Spec フィールド	推論関数	意味
distill	distill.pretrained_teacher_model_path	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
distill	results_dir	output_dir	現在のジョブの結果ディレクトリ
evaluate	evaluate.checkpoint	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
evaluate	results_dir	output_dir	現在のジョブの結果ディレクトリ
export	export.checkpoint	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
export	export.onnx_file	create_onnx_file	出力 ONNX パス
export	results_dir	output_dir	現在のジョブの結果ディレクトリ
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	inference.checkpoint	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
inference	inference.trt_engine	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
inference	results_dir	output_dir	現在のジョブの結果ディレクトリ
quantize	quantize.model_path	parent_model	親ジョブの結果フォルダーから推論されるモデルファイル
quantize	results_dir	output_dir	現在のジョブの結果ディレクトリ
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-image-classification — TAO Deploy を使用して TensorRT エンジンの生成、TensorRT による評価、TensorRT による推論を行う、Classification PyT のデプロイワークフロー。