tao-train-ocdnet

OCDNet

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

设置 train.pretrained_model_path 以指定预训练权重。

对于 TAO Deploy TensorRT 操作（gen_trt_engine、TensorRT evaluate 和 TensorRT inference），请先阅读 references/tao-deploy-ocdnet.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、默认值、最小值/最大值边界、枚举、选项权重、数学条件、依赖项和常用参数，请使用打包的训练架构。运行时不要依赖 ~/tao-core；维护者会在打包技能库之前重新生成架构和模板。

训练操作策略

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

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

训练要求

• 数据集类型： ocdnet
• 格式： 默认
• 监控指标： hmean

各操作的数据集要求

操作	规范键	来源	文件	列表？
evaluate	dataset.validate_dataset.data_path	eval_dataset	test.tar.gz	是
gen_trt_engine	gen_trt_engine.tensorrt.calibration.cal_image_dir	calibration_dataset	train/img.tar.gz	是
inference	inference.input_folder	eval_dataset	test/img.tar.gz	否
prune	dataset.validate_dataset.data_path	eval_dataset	test.tar.gz	是
quantize	dataset.train_dataset.data_path	train_datasets	train.tar.gz	是
quantize	dataset.validate_dataset.data_path	eval_dataset	test.tar.gz	是
quantize	dataset.quant_calibration_dataset.images_dir	train_datasets	train/img.tar.gz	否
retrain	dataset.train_dataset.data_path	train_datasets	train.tar.gz	是
retrain	dataset.validate_dataset.data_path	eval_dataset	test.tar.gz	是
train	dataset.train_dataset.data_path	train_datasets	train.tar.gz	是
train	dataset.validate_dataset.data_path	eval_dataset	test.tar.gz	是

典型规格覆盖项

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

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

train（必需数据源）：

{
    "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": [f"{S3_TRAIN}/train.tar.gz"],
    "dataset.validate_dataset.data_path": [f"{S3_EVAL}/test.tar.gz"],
}

gen_trt_engine（必需数据源）：

{
    "gen_trt_engine.tensorrt.data_type": "INT8",
    "gen_trt_engine.tensorrt.calibration.cal_image_dir": [f"{S3_TRAIN}/train/img.tar.gz"],
}

evaluate（必需数据源）：

{
    "dataset.validate_dataset.data_path": [f"{S3_EVAL}/test.tar.gz"],
}

inference（必需数据源）：

{
    "inference.input_folder": f"{S3_EVAL}/test/img.tar.gz",
}

prune（必需数据源）：

{
    "dataset.validate_dataset.data_path": [f"{S3_EVAL}/test.tar.gz"],
}

quantize（必需数据源）：

{
    "dataset.train_dataset.data_path": [f"{S3_TRAIN}/train.tar.gz"],
    "dataset.validate_dataset.data_path": [f"{S3_EVAL}/test.tar.gz"],
    "dataset.quant_calibration_dataset.images_dir": f"{S3_TRAIN}/train/img.tar.gz",
}

retrain（必需数据源）：

{
    "dataset.train_dataset.data_path": [f"{S3_TRAIN}/train.tar.gz"],
    "dataset.validate_dataset.data_path": [f"{S3_EVAL}/test.tar.gz"],
}

评估数据集

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

重要参数

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

多 GPU / 多节点

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

规格键	描述	默认值
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。默认阈值对某些数据集可能过于严格。

规范参数／父模型推断

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

来自 TAO Core ocdnet.config.json 的推断映射：

操作	规范字段	推断函数	含义
评估	evaluate.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
评估	evaluate.trt_engine	parent_model	从父作业结果文件夹推断出的模型文件
评估	model.pruned_graph_path	pruned_model	父作业的剪枝模型
评估	results_dir	output_dir	当前作业结果目录
导出	export.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
导出	export.onnx_file	create_onnx_file	输出 ONNX 路径
导出	results_dir	output_dir	当前作业结果目录
生成 TensorRT 引擎	gen_trt_engine.onnx_file	parent_model	从父作业结果文件夹推断出的模型文件
生成 TensorRT 引擎	gen_trt_engine.tensorrt.calibration.cal_cache_file	create_cal_cache	校准缓存路径
生成 TensorRT 引擎	gen_trt_engine.trt_engine	create_engine_file	输出 TensorRT 引擎路径
生成 TensorRT 引擎	results_dir	output_dir	当前作业结果目录
推理	inference.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
推理	inference.trt_engine	parent_model	从父作业结果文件夹推断出的模型文件
推理	model.pruned_graph_path	pruned_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 — 使用 TAO Deploy 生成 TensorRT 引擎、执行 TensorRT 评估和 TensorRT 推理的 OCDNet 部署工作流。