tao-train-image-classification

PyT 分类

可独立安装？ 如果此会话并非由 TAO skill bank 插件初始化，请先运行 tao-setup 技能（主机预检、凭据、跨技能发现）。

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。

数据类架构

生成的 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；维护者会在打包技能库之前重新生成架构/模板。

训练操作策略

此模型在模型层启用了 AutoML。在处理任何训练阶段请求之前，请阅读 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 时，默认通过 tao-skill-bank:tao-run-automl 执行训练操作，并传入此模型的 skill_dir。保留工作流/应用程序对数据集、规范、输出目录、GPU/平台设置、父检查点和 automl_policy 的覆盖设置。仅当 automl_policy: off 或缺少已打包的训练架构/模板时，才直接训练模型；如果缺少架构，请说明此模型已启用 AutoML，但在生成架构之前无法运行 AutoML。

evaluate、inference、export 和部署流程等非训练操作仍由此模型技能处理。每次运行的 automl_policy 覆盖设置不会更改模型元数据。

训练要求

• 数据集类型： image_classification
• 格式： classification_pyt
• 监控指标： val_acc_1
• AutoML 训练范围： 普通 train 搜索中排除 distill.teacher.* 参数。这些字段属于单独的 distill 操作；尽管它们出现在合并生成的架构中，也不会改变 classification_pyt train 的行为。

各操作的数据集要求

操作	规格键	来源	文件	列表？
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	否

常见规格覆盖项

每个操作都必须覆盖数据源——Agent 必须根据上方的各操作数据集要求表构建数据源路径，并将其包含在 spec_overrides 中。

TRAIN_IMAGES_DIR = "/workspace/data/extracted/train/images_train"
VAL_IMAGES_DIR = "/workspace/data/extracted/val/images_val"
TEST_IMAGES_DIR = "/workspace/data/extracted/test/images_test"
CLASSES_FILE = "/workspace/data/s3/classes.txt"
WRITABLE_RESULTS_DIR = "/results"   # must be a writable bind, not the image CWD

对于本地 Docker，请下载 S3 归档文件并先解压，然后将 dataset.*.images_dir 指向解压后的类别根文件夹。不要直接将 images_train.tar.gz、images_val.tar.gz 或 images_test.tar.gz 传给 本地 Docker 规格；技能元数据将这些输入声明为文件夹。

train（必须指定数据源）：

{
    "train.num_epochs": 2,
    "train.validation_interval": 2,
    "train.checkpoint_interval": 2,
    "train.num_gpus": 1,
    "dataset.train_dataset.images_dir": TRAIN_IMAGES_DIR,
    "dataset.classes_file": CLASSES_FILE,
    "dataset.val_dataset.images_dir": VAL_IMAGES_DIR,
    "dataset.root_dir": WRITABLE_RESULTS_DIR,
}

dataset.root_dir 对 train 是必填项，不只是对 export 必填。CLDataset 会在设置期间将 classes.txt 写入 root_dir，因此它必须指向可写的绑定挂载点（结果挂载点是自然的选择）。规格模板默认将其设为 ''，这会使写入操作落入容器工作目录——在以 root 身份运行的容器中可以成功，但在任何降权为非 root 用户的运行器中都会失败：

PermissionError: [Errno 13] Permission denied: 'classes.txt'
  .../classification_pyt/dataloader/dataset.py, in CLDataset.__init__

只要存在可写的结果绑定挂载，本地 Docker 启动器就会使用主机 UID:GID，因此如果未设置 dataset.root_dir，非 root 训练就会遇到此问题。

export（必须指定数据源）：

{
    "export.input_height": 224,
    "export.input_width": 224,
    "dataset.root_dir": "/workspace/data/extracted",
}

gen_trt_engine：

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

inference（必须指定数据源）：

{
    "dataset.batch_size": 1,
    "dataset.val_dataset.images_dir": VAL_IMAGES_DIR,
    "dataset.classes_file": CLASSES_FILE,
    "dataset.test_dataset.images_dir": TEST_IMAGES_DIR,
}

distill（必须指定数据源）：

{
    "dataset.train_dataset.images_dir": TRAIN_IMAGES_DIR,
    "dataset.classes_file": CLASSES_FILE,
    "dataset.val_dataset.images_dir": VAL_IMAGES_DIR,
    "train.optim.policy": "step",
}

evaluate（必须指定数据源）：

{
    "dataset.val_dataset.images_dir": VAL_IMAGES_DIR,
    "dataset.classes_file": CLASSES_FILE,
    "dataset.test_dataset.images_dir": TEST_IMAGES_DIR,
}

quantize（必须指定数据源）：

{
    "dataset.train_dataset.images_dir": TRAIN_IMAGES_DIR,
    "dataset.classes_file": CLASSES_FILE,
    "dataset.val_dataset.images_dir": VAL_IMAGES_DIR,
    "dataset.quant_calibration_dataset.images_dir": TRAIN_IMAGES_DIR,
}

评估数据集

可选。验证图像与训练图像一起作为单独的 tar 文件提供。 对于未提供单独 images_test.tar.gz 的小型冒烟测试数据集，请将 dataset.test_dataset.images_dir 设为验证归档文件，以便 evaluate 和 inference 仍能测试检查点交接。

重要参数

• 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 Key	描述	默认值
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 或更大的 VRAM（V100 或 A100）。分类任务通常较轻量。使用 batch_size=8 时，大多数骨干网络在 224x224 分辨率下都能很好地运行于 16GB GPU。

错误情况

CUDA 内存不足：减小 batch_size 或使用更小的骨干网络。

num_classes 不匹配：确保 dataset.num_classes 与图像 tar 包中的实际类别目录数量以及 classes.txt 中的类别数量一致。

类别目录为空：classes.txt 中的每个类别都必须在对应子目录中至少有一张图像。

蒸馏调度器默认值：随附的蒸馏模板和 schema 使用 train.optim.policy: step。除非更新容器实现，否则蒸馏配置请保留该设置；7.0 PyT 蒸馏器不会为 train.optim.policy: linear 分配调度器间隔。

检查点交接：训练会生成 model_epoch_*.pth 检查点和 classifier_model_latest.pth 符号链接。对于 evaluate、inference、export、quantize、 distill 和 resume，请通过 SDK 解析器选择确切的目标 epoch 检查点；仅当用户明确要求使用最新版本时，才使用 latest 符号链接。

Spec 参数／父模型推断

特定模型的推断映射应放在此 MD 文件中，而不是 config.json 中。生成的运行器应读取本节，并在调用 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