⚡

OCDNet场景文本检测Skill tao-train-ocdnet

OCDNet场景文本检测技能，基于可微分二值化方法检测自然图像中的任意方向文本区域。支持训练、评估、导出、剪枝、量化、重新训练与推理，并提供TAO部署流程。适用于场景文字识别、OCR预处理、自动驾驶路牌检测等。关键词：OCDNet, 场景文本检测, 可微分二值化, 任意方向文本, 自然图像, TAO, 模型训练, 模型推理, TensorRT

视觉模型训练 0 次安装 0 次浏览 更新于 9/6/2026
名称 tao-train-ocdnet
描述 OCDNet用于场景文本检测。使用可微分二值化方法检测自然图像中的任意方向文本区域。当需要对TAO OCDNet模型进行训练、评估、导出、剪枝、量化、重新训练或推理时使用。触发短语包括“train OCDNet”、“scene text detection”、“arbitrary-oriented text boxes”、“differentiable binarization detector”。
开源协议 Apache-2.0 compatibility: 需要docker + nvidia-container-toolkit。 metadata:
版本 “0.1.0”
作者 NVIDIA Corporation allowed-tools: Read Bash tags: - text - detection

OCDNet

独立安装？ 如果本会话未由TAO技能库插件初始化，请先运行tao-setup技能（主机预检、凭证、跨技能发现）。

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

设置model.pretrained_model_path以加载预训练权重。

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

PyT OCDNet命令行支持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引擎生成由部署工作流负责。

数据类模式

生成的TAO核心模式打包在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训练，使用train_loss_epoch或train_loss作为优化指标，方向为minimize。Lightning进度日志会输出train_loss_epoch，TAO status.json在train_loss下记录相同的最终值。对于单轮本地AutoML冒烟测试，设置train.lr_scheduler.args.warmup_epoch: 0；如果warmup等于epoch预算，则训练器会在推荐报告指标前失败。非训练操作（如evaluate、inference、export和部署流）保持在此模型技能中。每次运行的automl_policy覆盖不会改变模型元数据。

训练要求

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

每个操作的数据集要求

操作 规范键 来源 运行值 列表？
evaluate dataset.validate_dataset.data_path eval_dataset 提取的验证集拆分文件夹，包含img/和gt/ 是
inference inference.input_folder inference_dataset或eval_dataset 提取的图像文件夹 否
prune dataset.validate_dataset.data_path eval_dataset 提取的验证集拆分文件夹，包含img/和gt/ 是
quantize dataset.train_dataset.data_path train_datasets 提取的训练集拆分文件夹，包含img/和gt/ 是
quantize dataset.validate_dataset.data_path eval_dataset 提取的验证集拆分文件夹，包含img/和gt/ 是
quantize dataset.quant_calibration_dataset.images_dir train_datasets或calibration_dataset 提取的校准图像文件夹 否
train dataset.train_dataset.data_path train_datasets 提取的训练集拆分文件夹，包含img/和gt/ 是
train dataset.validate_dataset.data_path eval_dataset 提取的验证集拆分文件夹，包含img/和gt/ 是
retrain dataset.train_dataset.data_path train_datasets 提取的训练集拆分文件夹，包含img/和gt/ 是
retrain dataset.validate_dataset.data_path eval_dataset 提取的验证集拆分文件夹，包含img/和gt/ 是

典型规范覆盖

数据源覆盖对于每个操作都是强制性的——代理必须根据上面的表格构建数据源路径，并将其包含在spec_overrides中。OCDNet在运行时不会解压数据集归档。如果源是train.tar.gz、test.tar.gz或img.tar.gz，请先解压并将拆分文件夹或图像文件夹传入规范。拆分文件夹必须包含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],
}

默认规范输出：

{
    "results_dir": "<可写输出目录>",
}

评估数据集

可选。测试数据集以单独tar包提供。

重要参数

  • model.backbone：默认deformable_resnet18。可变形卷积能改善不规则文本区域的检测。
  • train.optimizer.args.lr：学习率，默认0.001（Adam）。
  • postprocess.thresh：文本区域提取的二值化阈值。
  • postprocess.box_thresh：过滤检测框的置信度阈值。

多GPU / 多节点

启动方法：Lightning管理（单个python进程，Lightning生成worker）。

规范键 描述 默认值
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。
  • 将归档路径传给数据集：dataset.*.data_path不是归档路径。直接传train.tar.gz会导致dataloader将其当作UTF-8数据列表打开。应解压后传包含img/与gt/的拆分文件夹。
  • 量化检查点类型：不要将model_best.pth传给PyTorch量化路径。它可能缺少完整Lightning元数据。应使用对应epoch的完整检查点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 epoch检查点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。剪枝会产出pruned_<ch_sparsity>.pth；从剪枝图重新训练时，将确切剪枝产物用作model.pruned_graph_path。仅当用户明确要求最新时才使用最新检查点。

若重试PyTorch量化，请使用与目标epoch对应的完整model_epoch_<epoch>_step_<step>.pth。若用modelopt.onnx重试，将导出的ONNX传给quantize.model_path并确认运行时镜像包含modelopt.onnx.quantization。

规范参数 / 父模型推断

模型特有的推断映射放在此MD中而非config.json。生成的runner应在create_job()之前读取本段并使用SDK助手应用映射。这类似旧微服务的infer_params.py。

操作 规范字段 推断函数 含义
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 当前任务结果目录
inference inference.checkpoint parent_model 模型文件从父任务结果文件夹推断
inference results_dir output_dir 当前任务结果目录
prune prune.checkpoint parent_model 模型文件从父任务结果文件夹推断
prune results_dir output_dir 当前任务结果目录
quantize quantize.model_path parent_model 模型文件从父任务结果文件夹推断
quantize results_dir output_dir 当前任务结果目录
从剪枝重新训练 model.pruned_graph_path parent_model 从父级剪枝结果文件夹推断的精确剪枝模型文件
从剪枝重新训练 results_dir output_dir 当前任务结果目录
训练 model.pretrained_model_path ptm_if_no_resume_model 无恢复检查点时使用的预训练模型
训练 results_dir output_dir 当前任务结果目录
训练 train.resume_training_checkpoint_path resume_model 从当前任务结果文件夹推断的模型文件

对于parent_model或parent_model_folder，请将上游train/export/AutoML子任务ID作为parent_job_id传入。SDK会列出父结果文件夹、筛选检查点工件，并返回所选模型文件或文件夹。请勿将这些映射重新写回config.json，也不要修补生成的runner脚本去猜测检查点路径。

部署