tao-train-ocdnet
OCDNet for scene text detection. Detects arbitrary-oriented text regions in natural images using a differentiable binarization approach. Use when training, evaluating, exporting, pruning, quantizing, retraining, or running inference for a TAO OCDNet model. Trigger phrases include "train OCDNet", "sc
By nvidia · 1,546 installs
npx skills add nvidia/skills --skill tao-train-ocdnet
Source repository · Upstream listing
OCDNet
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).
OCDNet for scene text detection. Detects arbitrary oriented text regions in natural images using a differentiable binarization approach.
Set model.pretrained model path for pretrained weights.
For TAO Deploy TensorRT actions ( gen trt engine , TensorRT evaluate , and TensorRT inference ), read references/tao deploy ocdnet.md first. Deploy spec templates live in this skill's references/ folder with the spec template deploy .yaml prefix.
The PyT OCDNet CLI supports train , evaluate , export , inference , prune , quantize , and default specs . It does not expose PyT side retrain or gen trt engine subcommands. The model skill exposes retrain by running ocdnet train with model.load pruned graph: true and model.pruned graph path . Resume from an epoch checkpoint uses ocdnet train plus train.resume training checkpoint path . TensorRT engine generation is owned by the deploy workflow.
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.
For AutoML train, use train loss epoch or train loss as the optimization
metric with direction=minimize . The Lightning progress log emits
train loss epoch , and TAO status.json records the same final value under
train loss . For one epoch local AutoML smoke runs, set
train.lr scheduler.args.warmup epoch: 0 ; leaving warmup equal to the epoch
budget causes the trainer to fail before a recommendation can report a metric.
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: ocdnet
Formats: default
Monitoring metric: hmean
Per Action Dataset Requirements
Action Spec Key Source Runtime value List?
evaluate dataset.validate dataset.data path eval dataset extracted validation split folder with img/ and gt/ Yes
inference inference.input folder inference dataset or eval dataset extracted image folder No
prune dataset.validate dataset.data path eval dataset extracted validation split folder with img/ and gt/ Yes
quantize dataset.train dataset.data path train datasets extracted train split folder with img/ and gt/ Yes
quantize dataset.validate dataset.data path eval dataset extracted validation split folder with img/ and gt/ Yes
quantize dataset.quant calibration dataset.images dir train datasets or calibration dataset extracted calibration image folder No
train dataset.train dataset.data path train datasets extracted train split folder with img/ and gt/ Yes
train dataset.validate dataset.data path eval dataset extracted validation split folder with img/ and gt/ Yes
retrain dataset.train dataset.data path train datasets extracted train split folder with img/ and gt/ Yes
retrain dataset.validate dataset.data path eval dataset extracted validation split folder with img/ and gt/ Yes
Typical Spec Overrides
Data source overrides are mandatory for every action — the agent MUST construct data source paths from the Per Action Dataset Requirements table above and include them in spec overrides . OCDNet does not unpack dataset archives at runtime. If the source is train.tar.gz , test.tar.gz , or img.tar.gz , extract it first and pass the split folder or image folder into the spec. The split folder must contain img/ and gt/ ; alternatively, pass a UTF 8 datalist text file whose lines map image paths to label paths.
train (mandatory data sources):
evaluate (mandatory data sources):
inference (mandatory data sources):
prune (mandatory data sources):
quantize (mandatory data sources):
resume training (mandatory data sources):
retrain from prune output (mandatory data sources):
default specs:
Eval Dataset
Optional. Test dataset provided as separate tarball.
Important Parameters
model.backbone : Default deformable resnet18. Deformable convolutions improve text region detection for irregular text.
train.optimizer.args.lr : Learning rate. Default 0.001 (Adam).
postprocess.thresh : Binarization threshold for text region extraction.
postprocess.box thresh : Box confidence threshold for filtering detections.
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 ddp , fsdp , or deepspeed stage 3 offload ddp
ddp with activation checkpointing: find unused parameters=False
ddp without: find unused parameters=True
fsdp forces FP16
deepspeed stage 3 offload is uniquely supported for OCDNet (forces FP16)
FAN backbones auto enable sync batchnorm
Hardware
Minimum 1 GPU(s), recommended 1 GPU(s). 8GB+ VRAM per GPU. OCDNet is lightweight. Single GPU is sufficient for most datasets.
Error Patterns
Low detection rate : Tune postprocess.thresh and box thresh. Default thresholds may be too aggressive for some datasets.
One epoch smoke train with default scheduler : train.num epochs must not equal train.lr scheduler.args.warmup epoch . For one epoch validation, set warmup epoch: 0 ; for normal starter runs, keep num epochs warmup epoch .
Archive passed as dataset path : dataset. .data path is not an archive path for OCDNet. Passing train.tar.gz or test.tar.gz directly causes the dataloader to open the gzip as a UTF 8 datalist. Extract the archive and pass the split folder containing img/ and gt/ , or pass a real UTF 8 datalist file.
Quantize checkpoint type : Do not pass model best.pth to the PyTorch quantize path. Some older PyT runtimes wrote model best.pth without full Lightning checkpoint metadata. The default torchao quantize path should use the intended full model epoch <epoch step <step .pth checkpoint and write quantized model torchao.pth .
Default specs output directory : ocdnet default specs requires a writable results dir override, for example results dir=/workspace/run/results/default specs .
Checkpoint Handoff
OCDNet train writes model best.pth plus full Lightning epoch checkpoints such as model epoch 001 step 00046.pth ; it may also write ocd model latest.pth as a latest symlink. Use model best.pth for evaluate.checkpoint , inference.checkpoint , export.checkpoint , and prune.checkpoint when the user asks for the best checkpoint. Use a specific model epoch <epoch step <step .pth for train.resume training checkpoint path and for any action that explicitly needs a full Lightning checkpoint. Prune writes artifacts such as pruned <ch sparsity .pth ; use the exact pruned .pth artifact for model.pruned graph path when retraining from a pruned graph. Use a latest checkpoint only when the user explicitly asks for latest.
If quantize is retried with a PyTorch backend, resolve the full model epoch <epoch step <step .pth that corresponds to the intended best epoch or requested epoch; do not pass model best.pth to the PyTorch quantize path. If quantize is retried with modelopt.onnx , pass the exported ONNX as quantize.model path and verify that the runtime image actually contains modelopt.onnx.quantization .
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.
Model handoff mappings:
Action Spec Field Inference Function Meaning
evaluate evaluate.checkpoint parent model model file inferred from the parent job results folder
evaluate results dir output dir current job results directory
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
inference inference.checkpoint parent model model file inferred from the parent job results folder
inference results dir output dir current job results directory
prune prune.checkpoint parent model model file inferred from the parent job results folder
prune results dir output dir current job results directory
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 from prune model.pruned graph path parent model exact pruned model file inferred from the parent prune results folder
retrain from prune results dir output dir current job results directory
train model.pretrained model path ptm if no resume model PTM when no resume checkpoint exists
train results dir output dir current job results directory
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 ocdnet](references/tao deploy ocdnet.md)