tao-train-ocrnet

--- name: tao-train-ocrnet description: OCRNet for scene text recognition. Recognizes text content from cropped text-region images and supports CTC and attention-based decoders. Use when training, evaluating, exporting, pruning, quantizing, retraining, or running inference for a TAO OCRNet model. Trigger phrases include "train OCRNet", "scene text recognition", "OCR cropped text", "CTC / attention text decoder". license: Apache-2.0 compatibility: Requires docker + nvidia-container-toolkit. metadata: version: "0.1.0" author: NVIDIA Corporation allowed-tools: Read Bash tags: - text - recognition --- # OCRNet > **Standalone install?** If this session was not initialized by the TAO skill bank plugin, run the `tao-setup` skill first (host preflight, credentials, cross-skill discovery). OCRNet for scene text recognition. Recognizes text content from cropped text region images. Supports CTC and attention-based decoders. Set train.pretrained_model_path for pretrained OCR weights. For TAO Deploy TensorRT actions (`gen_trt_engine`, TensorRT `evaluate`, and TensorRT `inference`), read `references/tao-deploy-ocrnet.md` first. Deploy spec templates live in this skill's `references/` folder with the `spec_template_deploy_*.yaml` prefix. ## Dataclass Schemas Generated TAO Core schemas are packaged in `schemas/<action>.schema.json`, with `schemas/manifest.json` listing available actions. Each generated schema also emits `references/spec_template_<action>.yaml` from the schema top-level `default` field. AutoML enablement is declared at the model layer in `references/skill_info.yaml` via `automl_enabled`. Runnable AutoML for an action requires `schemas/<action>.schema.json` and `references/spec_template_<action>.yaml` to exist and parse. Use the packaged selected-action schema for `automl_default_parameters`, `automl_disabled_parameters`, defaults, min/max bounds, enums, option weights, math conditions, dependencies, and popular parameters. Do not expect `~/tao-core` at runtime; maintainers regenerate schemas/templates before packaging the skill bank. ## Train Action Policy This model is AutoML-enabled at the model layer. Before handling any train-stage request, read `references/skill_info.yaml` and resolve the run override from either an explicit `automl_policy` value or the user's workflow request. Use `automl_policy: on` by default and only expose `on` / `off` in new launch prompts. Treat phrases like "turn off AutoML", "disable AutoML", "no HPO", or "plain training" as `automl_policy: off` for this run only. When `automl_policy: on`, `automl_enabled: true`, and both `schemas/train.schema.json` and `references/spec_template_train.yaml` are packaged, route the train action through `tao-skill-bank:tao-run-automl` by default with this model's `skill_dir`. Preserve workflow/application overrides for datasets, specs, output directories, GPU/platform settings, parent checkpoints, and `automl_policy`. Use direct model training only when `automl_policy: off` or the packaged train schema/template is missing; in the missing-schema case, report that AutoML is enabled but not runnable for this model until schemas are generated. Non-train actions such as `evaluate`, `inference`, `export`, and deploy flows stay in this model skill. The per-run `automl_policy` override does not change model metadata. ## Training Requirements - **Dataset type:** ocrnet - **Formats:** default - **AutoML training metric:** `val_acc`, with `direction=maximize`; TAO writes this key to `status.json` - **Lightning progress alias:** `val_acc_1` - **Standalone evaluation metric:** `test_acc`, with `direction=maximize` ### Per-Action Dataset Requirements | Action | Spec Key | Source | Files | List? | |---|---|---|---|---| | dataset_convert | dataset_convert.input_img_dir | train_datasets or eval_dataset | extracted folder containing cropped text images | No | | dataset_convert | dataset_convert.gt_file | train_datasets or eval_dataset | train/gt_new.txt or test/gt_new.txt | No | | evaluate | dataset.character_list_file | eval_dataset | character_list | No | | evaluate | evaluate.test_dataset_dir | eval_dataset | extracted test image folder | No | | evaluate | evaluate.test_dataset_gt_file | eval_dataset | test/gt_new.txt | No | | evaluate | evaluate.checkpoint | parent train/AutoML job | best_accuracy.pth or exact requested epoch checkpoint | No | | export | dataset.character_list_file | eval_dataset | character_list | No | | export | export.checkpoint | parent train/AutoML job | best_accuracy.pth or exact requested epoch checkpoint | No | | deploy/gen_trt_engine | gen_trt_engine.tensorrt.calibration.cal_image_dir | calibration_dataset | extracted calibration image folder for INT8 calibration | Yes | | deploy/gen_trt_engine | gen_trt_engine.onnx_file | parent export job | exported .onnx artifact | No | | deploy/gen_trt_engine | dataset.character_list_file | eval_dataset | character_list | No | | inference | dataset.character_list_file | eval_dataset | character_list | No | | inference | inference.inference_dataset_dir | inference_dataset | extracted inference image folder | No | | inference | inference.checkpoint | parent train/AutoML job | best_accuracy.pth or exact requested epoch checkpoint | No | | prune | dataset.character_list_file | eval_dataset | character_list | No | | prune | prune.checkpoint | parent train/AutoML job | best_accuracy.pth or exact requested epoch checkpoint | No | | quantize | dataset.train_dataset_dir | dataset_convert train job | LMDB folder containing data.mdb and lock.mdb | Yes | | quantize | dataset.val_dataset_dir | dataset_convert eval job | LMDB folder containing data.mdb and lock.mdb | No | | quantize | dataset.character_list_file | eval_dataset | character_list | No | | quantize | dataset.quant_calibration_dataset.images_dir | train_datasets | extracted calibration image folder | No | | quantize | quantize.model_path | parent train/AutoML job | checkpoint selected by resolver | No | | retrain | dataset.train_dataset_dir | dataset_convert train job | LMDB folder containing data.mdb and lock.mdb | Yes | | retrain | dataset.val_dataset_dir | dataset_convert eval job | LMDB folder containing data.mdb and lock.mdb | No | | retrain | dataset.character_list_file | eval_dataset | character_list | No | | retrain | model.pruned_graph_path | parent prune job | pruned .pth artifact | No | | train | dataset.train_dataset_dir | dataset_convert train job | LMDB folder containing data.mdb and lock.mdb | Yes | | train | dataset.train_gt_file | train_datasets | train/gt_new.txt when using raw folders instead of LMDB | No | | train | dataset.val_dataset_dir | dataset_convert eval job | LMDB folder containing data.mdb and lock.mdb | No | | train | dataset.val_gt_file | eval_dataset | test/gt_new.txt when using raw folders instead of LMDB | No | | train | dataset.character_list_file | eval_dataset | character_list | No | ### Checkpoint Selection OCRNet training writes both `best_accuracy.pth` and epoch-step checkpoints such as `model_epoch_000_step_00003.pth`. Use the SDK/model checkpoint resolver through the `spec_params` mappings in `references/skill_info.yaml`; do not guess by sorting for the newest `.pth`. - Use `best_accuracy.pth` for best-checkpoint `evaluate`, `inference`, `export`, and `prune` requests. - Use the exact requested `model_epoch_*_step_*.pth` for epoch/step-specific actions. - Use `train.resume_training_checkpoint_path` only for resume training, and use `model.pruned_graph_path` for retrain from a prune output. OCRNet does not expose a separate `ocrnet retrain` CLI subtask in the PyT image; the model-skill `retrain` action routes through `ocrnet train -e` with the pruned graph path set. - OCRNet `quantize` loads the model through PyTorch. For trusted checkpoints created by the same local run, set `TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1` if PyTorch 2.6+ rejects the checkpoint as a weights-only load. ### Typical Spec Overrides Data source overrides are **mandatory for every action**. Run `dataset_convert` separately for train and validation splits, then pass the LMDB folders that directly contain `data.mdb` and `lock.mdb` into train, quantize, and retrain. Tarballs from remote storage must be extracted before they are used as image directories. ```python 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 (run once per split):** ```python { "dataset_convert.input_img_dir": TRAIN_IMAGES, "dataset_convert.gt_file": TRAIN_GT, } ``` **train (mandatory data sources):** ```python { "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 (mandatory data sources):** ```python { "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 (mandatory data sources):** ```python { "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 (mandatory data sources):** ```python { "export.checkpoint": "<selected train/AutoML checkpoint>", "export.onnx_file": "<output ONNX path>", "dataset.character_list_file": CHAR_LIST, } ``` **inference (mandatory data sources):** ```python { "inference.checkpoint": "<selected train/AutoML checkpoint>", "dataset.character_list_file": CHAR_LIST, "inference.inference_dataset_dir": EVAL_IMAGES, } ``` **prune (mandatory data sources):** ```python { "prune.checkpoint": "<selected train/AutoML checkpoint>", "prune.pruned_file": "<output pruned PTH path>", "dataset.character_list_file": CHAR_LIST, } ``` **quantize (mandatory data sources):** ```python { "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 (mandatory data sources):** ```python { "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>", } ``` ## Eval Dataset Optional. Test data provided as separate tarball. ## Important Parameters - **dataset.character_list_file**: Path to character list defining the supported character set. This determines the output vocabulary size. - **model.backbone**: Default ResNet. - **model.prediction**: Decoder type. CTC or Attn (attention-based). - **train.optim.lr**: Learning rate. Default 1.0 (Adadelta optimizer). High default is specific to Adadelta. - **dataset.batch_size**: Per-GPU batch size. Default 16. ## Multi-GPU / Multi-Node **Launch method:** Lightning-managed (single `python` process, Lightning spawns workers). | Spec Key | Description | Default | |----------|-------------|---------| | `train.num_gpus` | Number of GPUs | 1 | | `train.gpu_ids` | GPU device indices | [0] | | `train.distributed_strategy` | Strategy name | `auto` | - Strategy: `auto` for single-GPU, reads `train.distributed_strategy` from config when multi-GPU - No explicit `num_nodes` in train script — single-node oriented - Lightweight model, single GPU typically sufficient ## Hardware Minimum 1 GPU(s), recommended 1 GPU(s). 8GB+ VRAM per GPU. OCR text recognition is lightweight. Single GPU is typically sufficient. ## Error Patterns **dataset_convert required**: If using raw images + gt files, run dataset_convert first to produce LMDB format. **dataset_convert output folder**: Direct `ocrnet dataset_convert` writes `data.mdb` and `lock.mdb` directly under `dataset_convert.results_dir`. Use that folder itself for `dataset.train_dataset_dir`, `dataset.val_dataset_dir`, quantize, and retrain inputs. SDK-backed runs may wrap the same LMDB folder inside job artifact directories; resolve the actual folder containing `data.mdb` and `lock.mdb`. **GT file BOM**: Some text-recognition GT files can start with a UTF-8 BOM on the first filename. If dataset conversion logs a missing path with an invisible prefix before the first image name, strip the BOM from a local copy of the GT file before conversion or evaluation. **Character list mismatch**: All characters in training data must be present in the character_list file. **Export/prune output fields required**: `export.onnx_file` and `prune.pruned_file` must be writable output paths. These are declared in `references/skill_info.yaml` so SDK-backed model runs can create the paths automatically. **TensorRT lives in deploy**: The PyT OCRNet CLI exposes `dataset_convert`, `evaluate`, `export`, `inference`, `prune`, `quantize`, and `train`, but not `gen_trt_engine`. Use `references/tao-deploy-ocrnet.md` and `deploy/skill_info.yaml` for TensorRT engine generation and TensorRT-backed evaluate/inference. ## Spec Param / Parent Model Inference Model-specific inference mappings belong in this MD file, not in `config.json`. Generated runners should read this section and apply the mappings with SDK helpers before `create_job()`. This mirrors the old microservices `infer_params.py` flow. Inference mappings from TAO Core `ocrnet.config.json`: | Action | Spec Field | Inference Function | Meaning | |---|---|---|---| | dataset_convert | `results_dir` | `output_dir` | current job results directory | | evaluate | `encryption_key` | `key` | encryption key | | evaluate | `evaluate.checkpoint` | `parent_model` | model file inferred from the parent job results folder | | evaluate | `evaluate.trt_engine` | `parent_model` | model file inferred from the parent job results folder | | evaluate | `model.pruned_graph_path` | `pruned_model` | parent pruned model | | evaluate | `results_dir` | `output_dir` | current job results directory | | export | `encryption_key` | `key` | encryption key | | export | `export.checkpoint` | `parent_model` | model file inferred from the parent job results folder | | export | `export.onnx_file` | `create_onnx_file` | output ONNX path | | export | `results_dir` | `output_dir` | current job results directory | | deploy/gen_trt_engine | `encryption_key` | `key` | encryption key | | deploy/gen_trt_engine | `gen_trt_engine.onnx_file` | `parent_model` | ONNX file inferred from the parent export job results folder | | deploy/gen_trt_engine | `gen_trt_engine.tensorrt.calibration.cal_cache_file` | `create_cal_cache` | calibration cache path | | deploy/gen_trt_engine | `gen_trt_engine.trt_engine` | `create_engine_file` | output TensorRT engine path | | deploy/gen_trt_engine | `results_dir` | `output_dir` | current job results directory | | inference | `encryption_key` | `key` | encryption key | | inference | `inference.checkpoint` | `parent_model` | model file inferred from the parent job results folder | | inference | `inference.trt_engine` | `parent_model` | model file inferred from the parent job results folder | | inference | `model.pruned_graph_path` | `pruned_model` | parent pruned model | | inference | `results_dir` | `output_dir` | current job results directory | | prune | `encryption_key` | `key` | encryption key | | prune | `prune.checkpoint` | `parent_model` | model file inferred from the parent job results folder | | prune | `prune.pruned_file` | `create_pth_file` | output PTH path | | prune | `results_dir` | `output_dir` | current job results directory | | quantize | `encryption_key` | `key` | encryption key | | quantize | `quantize.model_path` | `parent_model` | model file inferred from the parent job results folder | | quantize | `results_dir` | `output_dir` | current job results directory | | retrain | `encryption_key` | `key` | encryption key | | retrain | `model.pruned_graph_path` | `parent_model` | model file inferred from the parent job results folder | | retrain | `results_dir` | `output_dir` | current job results directory | | train | `encryption_key` | `key` | encryption key | | train | `results_dir` | `output_dir` | current job results directory | | train | `train.pretrained_model_path` | `ptm_if_no_resume_model` | PTM when no resume checkpoint exists | | train | `train.resume_training_checkpoint_path` | `resume_model` | model file inferred from the current job results folder | For `parent_model` or `parent_model_folder`, pass the upstream train/export/AutoML child job id as `parent_job_id`. The SDK lists the parent result folder, filters checkpoint artifacts, and returns the selected model file or folder. Do not add these mappings back to `config.json` and do not patch generated runner scripts to guess checkpoint paths. ## Deployment - [tao-deploy-ocrnet](references/tao-deploy-ocrnet.md)

存档时间线

版本索引永久保留；内容副本只保留最新 1 份，更早版本仅留索引与哈希（存档时间线的证据链不会因此断裂）。

下载存档内容副本

相关资产