tao-train-ocdnet

OCDNet

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

OCDNet 用于场景文本检测。它采用可微二值化方法检测自然图像中任意方向的文本区域。

为预训练权重设置 model.pretrained_model_path。

对于 TAO Deploy TensorRT 操作（gen_trt_engine、TensorRT evaluate 和 TensorRT inference），请先阅读 references/tao-deploy-ocdnet.md。Deploy 规格模板位于此技能的 references/ 文件夹中，前缀为 spec_template_deploy_*.yaml。

PyT OCDNet CLI 支持 train、evaluate、export、inference、prune、quantize 和 default_specs。它不提供 PyT 端的 retrain 或 gen_trt_engine 子命令。模型技能通过运行 ocdnet train，并设置 model.load_pruned_graph: true 和 model.pruned_graph_path 来提供 retrain。从 epoch 检查点恢复时，使用 ocdnet train 并设置 train.resume_training_checkpoint_path。TensorRT 引擎生成由 deploy 工作流负责。

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/<action>.schema.json 和 references/spec_template_<action>.yaml。对于 automl_default_parameters、automl_disabled_parameters、默认值、最小值/最大值范围、枚举、选项权重、数学条件、依赖项和常用参数，请使用打包的所选操作架构。运行时不要依赖 ~/tao-core；维护者会在打包 skill bank 之前重新生成架构/模板。

训练操作策略

此模型在模型层级启用了 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。

对于 AutoML 训练，使用 train_loss_epoch 或 train_loss 作为优化指标，并将 direction=minimize。Lightning 进度日志会输出 train_loss_epoch，TAO status.json 会在 train_loss 下记录相同的最终值。对于单 epoch 本地 AutoML 冒烟测试，请设置 train.lr_scheduler.args.warmup_epoch: 0；如果 warmup 与 epoch 预算相同，训练器会在生成建议并报告指标之前失败。evaluate、inference、export 和 deploy 流程等非训练操作仍由此模型技能处理。每次运行的 automl_policy 覆盖不会更改模型元数据。

训练要求

• 数据集类型： ocdnet
• 格式： default
• AutoML 训练指标： train_loss（状态记录器最终的 train_loss_epoch 值），其中 direction=minimize
• AutoML 指标约定： 使用训练期间输出的 train_loss，并将其最小化。仅将 train_loss_epoch 视为日志回退别名；仅使用独立的 hmean 验证所选检查点。
• 独立评估指标： hmean，其中 direction=maximize

各操作的数据集要求

操作	Spec Key	来源	运行时值	列表？
评估	dataset.validate_dataset.data_path	eval_dataset	已提取的验证集拆分文件夹，包含 img/ 和 gt/	是
推理	inference.input_folder	inference_dataset 或 eval_dataset	已提取的图像文件夹	否
剪枝	dataset.validate_dataset.data_path	eval_dataset	已提取的验证集拆分文件夹，包含 img/ 和 gt/	是
量化	dataset.train_dataset.data_path	train_datasets	已提取的训练集拆分文件夹，包含 img/ 和 gt/	是
量化	dataset.validate_dataset.data_path	eval_dataset	已提取的验证集拆分文件夹，包含 img/ 和 gt/	是
量化	dataset.quant_calibration_dataset.images_dir	train_datasets 或 calibration_dataset	已提取的校准图像文件夹	否
训练	dataset.train_dataset.data_path	train_datasets	已提取的训练集拆分文件夹，包含 img/ 和 gt/	是
训练	dataset.validate_dataset.data_path	eval_dataset	已提取的验证集拆分文件夹，包含 img/ 和 gt/	是
重新训练	dataset.train_dataset.data_path	train_datasets	已提取的训练集拆分文件夹，包含 img/ 和 gt/	是
重新训练	dataset.validate_dataset.data_path	eval_dataset	已提取的验证集拆分文件夹，包含 img/ 和 gt/	是

常见 Spec 覆盖项

每个操作都必须覆盖数据源——Agent 必须根据上方“各操作的数据集要求”表构建数据源路径，并将其包含在 spec_overrides 中。OCDNet 不会在运行时解压数据集归档文件。如果数据源是 train.tar.gz、test.tar.gz 或 img.tar.gz，请先解压，然后将拆分文件夹或图像文件夹传入 spec。拆分文件夹必须包含 img/ 和 gt/；或者，也可以传入 UTF-8 数据列表文本文件，其中每行将图像路径映射到标签路径。

TRAIN_ROOT = "/path/to/extracted/train"
EVAL_ROOT = "/path/to/extracted/test"
INFER_IMG_DIR = "/path/to/extracted/test/img"
CALIB_IMG_DIR = "/path/to/extracted/train/img"

训练（必须提供数据源）：

{
    "train.num_epochs": 30,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "train.num_gpus": 1,
    "dataset.train_dataset.loader.batch_size": 16,
    "dataset.train_dataset.data_path": [TRAIN_ROOT],
    "dataset.validate_dataset.data_path": [EVAL_ROOT],
}

评估（必须提供数据源）：

{
    "evaluate.checkpoint": "<selected train/AutoML checkpoint>",
    "dataset.validate_dataset.data_path": [EVAL_ROOT],
}

推理（必须提供数据源）：

{
    "inference.checkpoint": "<selected train/AutoML checkpoint>",
    "inference.input_folder": INFER_IMG_DIR,
}

剪枝（必须提供数据源）：

{
    "prune.checkpoint": "<selected train/AutoML checkpoint>",
    "dataset.validate_dataset.data_path": [EVAL_ROOT],
}

量化（必须提供数据源）：

{
    "quantize.model_path": "<selected train checkpoint or exported ONNX>",
    "dataset.train_dataset.data_path": [TRAIN_ROOT],
    "dataset.validate_dataset.data_path": [EVAL_ROOT],
    "dataset.quant_calibration_dataset.images_dir": CALIB_IMG_DIR,
}

恢复训练（必须提供数据源）：

{
    "train.resume_training_checkpoint_path": "<exact model_epoch checkpoint>",
    "dataset.train_dataset.data_path": [TRAIN_ROOT],
    "dataset.validate_dataset.data_path": [EVAL_ROOT],
}

从剪枝输出重新训练（必须提供数据源）：

{
    "model.load_pruned_graph": True,
    "model.pruned_graph_path": "<selected prune output>",
    "dataset.train_dataset.data_path": [TRAIN_ROOT],
    "dataset.validate_dataset.data_path": [EVAL_ROOT],
}

default_specs：

{
    "results_dir": "<writable output directory>",
}

评估数据集

可选。以单独 tarball 形式提供的测试数据集。

重要参数

• model.backbone：默认值为 deformable_resnet18。可变形卷积可提升对不规则文本的文本区域检测效果。
• train.optimizer.args.lr：学习率。默认值为 0.001（Adam）。
• postprocess.thresh：文本区域提取的二值化阈值。
• postprocess.box_thresh：用于过滤检测结果的边界框置信度阈值。

多 GPU / 多节点

启动方式： 由 Lightning 管理（单个 python 进程，Lightning 会启动工作进程）。

Spec Key	描述	默认值
train.num_gpus	GPU 数量	1
train.gpu_ids	GPU 设备索引	[0]
train.distributed_strategy	ddp、fsdp 或 deepspeed_stage_3_offload	ddp

• 使用激活检查点时的 ddp：find_unused_parameters=False
• 不使用激活检查点时的 ddp：find_unused_parameters=True
• fsdp 会强制使用 FP16
• deepspeed_stage_3_offload 是 OCDNet 独有支持的选项（会强制使用 FP16）
• FAN 主干网络会自动启用 sync_batchnorm

硬件

最低需要 1 块 GPU，推荐使用 1 块 GPU。每块 GPU 需有 8GB 以上显存。OCDNet 是轻量级模型。对于大多数数据集，单块 GPU 就足够。

错误模式

检测率低：调整 postprocess.thresh 和 box_thresh。默认阈值对某些数据集可能过于严格。

使用默认调度器进行单轮试运行：train.num_epochs 不得等于 train.lr_scheduler.args.warmup_epoch。进行单轮验证时，将 warmup_epoch: 0；进行常规的初始运行时，保持 num_epochs > warmup_epoch。

将归档文件作为数据集路径传入：对于 OCDNet，dataset.*.data_path 不能是归档文件路径。直接传入 train.tar.gz 或 test.tar.gz 会导致数据加载器将 gzip 文件作为 UTF-8 数据列表打开。请解压归档文件，并传入包含 img/ 和 gt/ 的数据拆分文件夹，或者传入真正的 UTF-8 数据列表文件。

量化检查点类型：不要将 model_best.pth 传入 PyTorch 量化路径。一些较旧的 PyT 运行时生成的 model_best.pth 不包含完整的 Lightning 检查点元数据。默认的 torchao 量化路径应使用预期的完整 model_epoch_<epoch>_step_<step>.pth 检查点，并写入 quantized_model_torchao.pth。

默认规格输出目录：ocdnet default_specs 需要可写的 results_dir 覆盖值，例如 results_dir=/workspace/run/results/default_specs。

检查点交接

OCDNet 训练会写入 model_best.pth 以及完整的 Lightning 轮次检查点，例如 model_epoch_001_step_00046.pth；也可能将 ocd_model_latest.pth 写为最新检查点的符号链接。用户要求使用最佳检查点时，请将 model_best.pth 用于 evaluate.checkpoint、inference.checkpoint、export.checkpoint 和 prune.checkpoint。对于 train.resume_training_checkpoint_path 以及任何明确需要完整 Lightning 检查点的操作，请使用指定的 model_epoch_<epoch>_step_<step>.pth。Prune 会写入 pruned_<ch_sparsity>.pth 等产物；从剪枝后的图重新训练时，请将准确的剪枝 .pth 产物用于 model.pruned_graph_path。仅当用户明确要求最新检查点时，才使用最新检查点。

如果使用 PyTorch 后端重试量化，请解析与预期的最佳轮次或指定轮次对应的完整 model_epoch_<epoch>_step_<step>.pth；不要将 model_best.pth 传入 PyTorch 量化路径。如果使用 modelopt.onnx 重试量化，请将导出的 ONNX 传入 quantize.model_path，并确认运行时镜像中确实包含 modelopt.onnx.quantization。

Spec 参数／父模型推断

特定于模型的推断映射应放在此 MD 文件中，而非 config.json。生成的运行器应读取本节，并在调用 create_job() 前使用 SDK 辅助函数应用这些映射。这与旧版微服务的 infer_params.py 流程一致。

模型传递映射：

操作	Spec 字段	推断函数	含义
评估	evaluate.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
评估	results_dir	output_dir	当前作业结果目录
导出	export.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
导出	export.onnx_file	create_onnx_file	输出 ONNX 路径
导出	results_dir	output_dir	当前作业结果目录
推理	inference.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
推理	results_dir	output_dir	当前作业结果目录
剪枝	prune.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
剪枝	results_dir	output_dir	当前作业结果目录
量化	quantize.model_path	parent_model	从父作业结果文件夹推断出的模型文件
量化	results_dir	output_dir	当前作业结果目录
从剪枝结果重新训练	model.pruned_graph_path	parent_model	从父剪枝作业结果文件夹推断出的确切剪枝模型文件
从剪枝结果重新训练	results_dir	output_dir	当前作业结果目录
训练	model.pretrained_model_path	ptm_if_no_resume_model	不存在恢复检查点时使用的 PTM
训练	results_dir	output_dir	当前作业结果目录
训练	train.resume_training_checkpoint_path	resume_model	从当前作业结果文件夹推断出的模型文件

对于 parent_model 或 parent_model_folder，请将上游训练／导出／AutoML 子作业 ID 作为 parent_job_id 传入。SDK 会列出父结果文件夹、筛选检查点工件，并返回选定的模型文件或文件夹。不要将这些映射加回 config.json，也不要修改生成的运行器脚本来猜测检查点路径。

部署

• tao-deploy-ocdnet