tao-train-ocrnet

OCRNet

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

用于场景文本识别的 OCRNet。可识别裁剪后的文本区域图像中的文本内容。支持 CTC 和基于注意力的解码器。

设置 train.pretrained_model_path 以使用预训练 OCR 权重。

对于 TAO Deploy TensorRT 操作（gen_trt_engine、TensorRT evaluate 和 TensorRT inference），请先阅读 references/tao-deploy-ocrnet.md。Deploy 规范模板位于此 skill 的 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/<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。

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

训练要求

• 数据集类型： ocrnet
• 格式： 默认
• AutoML 训练指标： val_acc，其中 direction=maximize；TAO 会将 此键写入 status.json
• Lightning 进度别名： val_acc_1
• 独立评估指标： test_acc，其中 direction=maximize

各操作的数据集要求

操作	规范键	来源	文件	列表？
dataset_convert	dataset_convert.input_img_dir	train_datasets 或 eval_dataset	包含裁剪文本图像的已解压文件夹	否
dataset_convert	dataset_convert.gt_file	train_datasets 或 eval_dataset	train/gt_new.txt 或 test/gt_new.txt	否
evaluate	dataset.character_list_file	eval_dataset	character_list	否
evaluate	evaluate.test_dataset_dir	eval_dataset	已解压的测试图像文件夹	否
evaluate	evaluate.test_dataset_gt_file	eval_dataset	test/gt_new.txt	否
evaluate	evaluate.checkpoint	父级 train/AutoML 作业	best_accuracy.pth 或指定的确切轮次检查点	否
export	dataset.character_list_file	eval_dataset	character_list	否
export	export.checkpoint	父级 train/AutoML 作业	best_accuracy.pth 或指定的确切轮次检查点	否
deploy/gen_trt_engine	gen_trt_engine.tensorrt.calibration.cal_image_dir	calibration_dataset	用于 INT8 校准的已解压校准图像文件夹	是
deploy/gen_trt_engine	gen_trt_engine.onnx_file	父级 export 作业	导出的 .onnx 产物	否
deploy/gen_trt_engine	dataset.character_list_file	eval_dataset	character_list	否
inference	dataset.character_list_file	eval_dataset	character_list	否
inference	inference.inference_dataset_dir	inference_dataset	已解压的推理图像文件夹	否
inference	inference.checkpoint	父级 train/AutoML 作业	best_accuracy.pth 或指定的确切轮次检查点	否
prune	dataset.character_list_file	eval_dataset	character_list	否
prune	prune.checkpoint	父级 train/AutoML 作业	best_accuracy.pth 或指定的确切轮次检查点	否
quantize	dataset.train_dataset_dir	dataset_convert train 作业	包含 data.mdb 和 lock.mdb 的 LMDB 文件夹	是
quantize	dataset.val_dataset_dir	dataset_convert eval 作业	包含 data.mdb 和 lock.mdb 的 LMDB 文件夹	否
quantize	dataset.character_list_file	eval_dataset	character_list	否
quantize	dataset.quant_calibration_dataset.images_dir	train_datasets	已解压的校准图像文件夹	否
quantize	quantize.model_path	父级 train/AutoML 作业	由解析器选定的检查点	否
retrain	dataset.train_dataset_dir	dataset_convert train 作业	包含 data.mdb 和 lock.mdb 的 LMDB 文件夹	是
retrain	dataset.val_dataset_dir	dataset_convert eval 作业	包含 data.mdb 和 lock.mdb 的 LMDB 文件夹	否
retrain	dataset.character_list_file	eval_dataset	character_list	否
retrain	model.pruned_graph_path	父级 prune 作业	剪枝后的 .pth 产物	否
train	dataset.train_dataset_dir	dataset_convert train 作业	包含 data.mdb 和 lock.mdb 的 LMDB 文件夹	是
train	dataset.train_gt_file	train_datasets	使用原始文件夹而非 LMDB 时的 train/gt_new.txt	否
train	dataset.val_dataset_dir	dataset_convert eval 作业	包含 data.mdb 和 lock.mdb 的 LMDB 文件夹	否
train	dataset.val_gt_file	eval_dataset	使用原始文件夹而非 LMDB 时的 test/gt_new.txt	否
train	dataset.character_list_file	eval_dataset	character_list	否

检查点选择

OCRNet 训练会同时生成 best_accuracy.pth 和按轮次、步数命名的检查点，例如 model_epoch_000_step_00003.pth。请通过 references/skill_info.yaml 中的 spec_params 映射使用 SDK／模型检查点解析器；不要猜测并按时间排序来选择最新的 .pth。

• 对最佳检查点相关的 evaluate、inference、export 和 prune 请求，使用 best_accuracy.pth。
• 对指定轮次／步数的操作，使用请求中指定的确切 model_epoch_*_step_*.pth。
• 仅在恢复训练时使用 train.resume_training_checkpoint_path；从剪枝输出重新训练时使用 model.pruned_graph_path。PyT 镜像中的 OCRNet 不提供单独的 ocrnet retrain CLI 子任务；模型技能的 retrain 操作会通过设置剪枝图路径来调用 ocrnet train -e。
• OCRNet quantize 会通过 PyTorch 加载模型。对于由同一次本地运行生成的可信检查点，如果 PyTorch 2.6+ 拒绝以仅加载权重的方式加载，请设置 TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1。

常见 Spec 覆盖项

每项操作都必须覆盖数据源。分别对训练集和验证集运行 dataset_convert，然后将直接包含 data.mdb 和 lock.mdb 的 LMDB 文件夹传入 train、quantize 和 retrain。远程存储中的压缩包必须先解压，才能用作图像目录。

TRAIN_IMAGES = "<extracted train image folder>"
TRAIN_GT = "<train gt_new.txt>"
EVAL_IMAGES = "<extracted eval image folder>"
EVAL_GT = "<eval gt_new.txt>"
TRAIN_LMDB = "<train dataset_convert results_dir>"
EVAL_LMDB = "<eval dataset_convert results_dir>"
CHAR_LIST = "<character_list>"

dataset_convert（每个数据集划分运行一次）：

{
    "dataset_convert.input_img_dir": TRAIN_IMAGES,
    "dataset_convert.gt_file": TRAIN_GT,
}

train（必须指定数据源）：

{
    "train.num_epochs": 30,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "train.num_gpus": 1,
    "dataset.batch_size": 16,
    "dataset.train_dataset_dir": [TRAIN_LMDB],
    "dataset.val_dataset_dir": EVAL_LMDB,
    "dataset.train_gt_file": "",
    "dataset.val_gt_file": "",
    "dataset.character_list_file": CHAR_LIST,
}

deploy/gen_trt_engine（必须指定数据源）：

{
    "gen_trt_engine.onnx_file": "<selected export ONNX>",
    "gen_trt_engine.trt_engine": "<output engine path>",
    "gen_trt_engine.tensorrt.calibration.cal_cache_file": "<output calibration cache path>",
    "gen_trt_engine.tensorrt.data_type": "fp16",
    "gen_trt_engine.tensorrt.calibration.cal_image_dir": [TRAIN_IMAGES],
    "dataset.character_list_file": CHAR_LIST,
}

evaluate（必须指定数据源）：

{
    "evaluate.checkpoint": "<selected train/AutoML checkpoint>",
    "dataset.character_list_file": CHAR_LIST,
    "evaluate.test_dataset_dir": EVAL_IMAGES,
    "evaluate.test_dataset_gt_file": EVAL_GT,
}

export（必须指定数据源）：

{
    "export.checkpoint": "<selected train/AutoML checkpoint>",
    "export.onnx_file": "<output ONNX path>",
    "dataset.character_list_file": CHAR_LIST,
}

inference（必须指定数据源）：

{
    "inference.checkpoint": "<selected train/AutoML checkpoint>",
    "dataset.character_list_file": CHAR_LIST,
    "inference.inference_dataset_dir": EVAL_IMAGES,
}

prune（必须指定数据源）：

{
    "prune.checkpoint": "<selected train/AutoML checkpoint>",
    "prune.pruned_file": "<output pruned PTH path>",
    "dataset.character_list_file": CHAR_LIST,
}

quantize（必须指定数据源）：

{
    "dataset.train_dataset_dir": [TRAIN_LMDB],
    "dataset.val_dataset_dir": EVAL_LMDB,
    "dataset.character_list_file": CHAR_LIST,
    "dataset.quant_calibration_dataset.images_dir": TRAIN_IMAGES,
    "quantize.model_path": "<selected train/AutoML checkpoint>",
}

retrain（必须指定数据源）：

{
    "dataset.train_dataset_dir": [TRAIN_LMDB],
    "dataset.val_dataset_dir": EVAL_LMDB,
    "dataset.character_list_file": CHAR_LIST,
    "model.pruned_graph_path": "<selected prune output>",
}

评估数据集

可选。测试数据以单独的压缩包形式提供。

重要参数

• dataset.character_list_file：定义支持字符集的字符列表路径。此参数决定输出词汇表大小。
• model.backbone：默认为 ResNet。
• model.prediction：解码器类型。CTC 或 Attn（基于注意力）。
• train.optim.lr：学习率。默认值为 1.0（Adadelta 优化器）。较高的默认值是 Adadelta 特有的。
• dataset.batch_size：每张 GPU 的批次大小。默认值为 16。

多 GPU／多节点

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

Spec Key	描述	默认值
train.num_gpus	GPU 数量	1
train.gpu_ids	GPU 设备索引	[0]
train.distributed_strategy	策略名称	auto

• 策略：单 GPU 时为 auto；多 GPU 时从配置中读取 train.distributed_strategy
• 训练脚本中未显式设置 num_nodes——面向单节点
• 模型轻量，通常单张 GPU 即可满足需求

硬件

最低需要 1 张 GPU，建议使用 1 张 GPU。每张 GPU 需配备 8GB 以上显存。OCR 文本识别的计算量较小，通常单张 GPU 即可满足需求。

错误模式

需要执行 dataset_convert：如果使用原始图像和 gt 文件，请先运行 dataset_convert，以生成 LMDB 格式的数据。

dataset_convert 输出文件夹：直接运行 ocrnet dataset_convert 会将 data.mdb 和 lock.mdb 直接写入 dataset_convert.results_dir。请将该文件夹本身用于 dataset.train_dataset_dir、dataset.val_dataset_dir、量化和重新训练输入。由 SDK 支持的运行可能会将同一个 LMDB 文件夹封装在作业产物目录中；请找到实际包含 data.mdb 和 lock.mdb 的文件夹。

GT 文件 BOM：某些文本识别 GT 文件的第一个文件名前可能带有 UTF-8 BOM。如果数据集转换日志显示缺失路径，且第一个图像名前有不可见前缀，请在转换或评估前，从 GT 文件的本地副本中去除 BOM。

字符列表不匹配：训练数据中的所有字符都必须包含在 character_list 文件中。

导出／剪枝输出字段必填：export.onnx_file 和 prune.pruned_file 必须是可写的输出路径。这些路径在 references/skill_info.yaml 中声明，以便由 SDK 支持的模型运行自动创建。

TensorRT 位于 deploy 中：PyT OCRNet CLI 提供 dataset_convert、evaluate、export、inference、prune、quantize 和 train，但不提供 gen_trt_engine。请使用 references/tao-deploy-ocrnet.md 和 deploy/skill_info.yaml 生成 TensorRT 引擎，以及执行基于 TensorRT 的评估／推理。

Spec 参数 / 父模型推断

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

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

操作	Spec 字段	推断函数	含义
dataset_convert	results_dir	output_dir	当前作业的结果目录
evaluate	encryption_key	key	加密密钥
evaluate	evaluate.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
evaluate	evaluate.trt_engine	parent_model	从父作业结果文件夹推断出的模型文件
evaluate	model.pruned_graph_path	pruned_model	父剪枝模型
evaluate	results_dir	output_dir	当前作业的结果目录
export	encryption_key	key	加密密钥
export	export.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
export	export.onnx_file	create_onnx_file	输出 ONNX 路径
export	results_dir	output_dir	当前作业的结果目录
deploy/gen_trt_engine	encryption_key	key	加密密钥
deploy/gen_trt_engine	gen_trt_engine.onnx_file	parent_model	从父导出作业结果文件夹推断出的 ONNX 文件
deploy/gen_trt_engine	gen_trt_engine.tensorrt.calibration.cal_cache_file	create_cal_cache	校准缓存路径
deploy/gen_trt_engine	gen_trt_engine.trt_engine	create_engine_file	输出 TensorRT 引擎路径
deploy/gen_trt_engine	results_dir	output_dir	当前作业的结果目录
inference	encryption_key	key	加密密钥
inference	inference.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
inference	inference.trt_engine	parent_model	从父作业结果文件夹推断出的模型文件
inference	model.pruned_graph_path	pruned_model	父剪枝模型
inference	results_dir	output_dir	当前作业的结果目录
prune	encryption_key	key	加密密钥
prune	prune.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
prune	prune.pruned_file	create_pth_file	输出 PTH 路径
prune	results_dir	output_dir	当前作业的结果目录
quantize	encryption_key	key	加密密钥
quantize	quantize.model_path	parent_model	从父作业结果文件夹推断出的模型文件
quantize	results_dir	output_dir	当前作业的结果目录
retrain	encryption_key	key	加密密钥
retrain	model.pruned_graph_path	parent_model	父剪枝模型
retrain	results_dir	output_dir	当前作业的结果目录
train	encryption_key	key	加密密钥
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-ocrnet