tao-train-ocrnet

Par nvidia · skills

OCRNet pour la reconnaissance de texte en scène. Reconnaît le contenu textuel à partir d'images recadrées de régions de texte et prend en charge le CTC

npx skills add https://github.com/nvidia/skills --skill tao-train-ocrnet

OCRNet

OCRNet pour la reconnaissance de texte en scène. Reconnaît le contenu texte à partir d'images de régions textuelles recadrées. Supporte les décodeurs CTC et basés sur l'attention.

Définissez train.pretrained_model_path pour les poids OCR préentraînés.

Pour les actions TAO Deploy TensorRT (gen_trt_engine, TensorRT evaluate, et TensorRT inference), lisez d'abord references/tao-deploy-ocrnet.md. Les modèles de spec de déploiement se trouvent dans le dossier references/ de cette skill avec le préfixe spec_template_deploy_*.yaml.

Schémas Dataclass

Les schémas TAO Core générés sont fournis dans schemas/<action>.schema.json, avec schemas/manifest.json listant les actions disponibles. Chaque schéma généré émet également references/spec_template_<action>.yaml à partir du champ default au niveau supérieur du schéma. L'activation AutoML est déclarée au niveau du modèle dans references/skill_info.yaml via automl_enabled. L'AutoML exécutable nécessite toujours que schemas/train.schema.json et references/spec_template_train.yaml existent et analysent correctement. Utilisez le schéma train fourni pour automl_default_parameters, automl_disabled_parameters, les valeurs par défaut, les limites min/max, les énumérations, les poids des options, les conditions mathématiques, les dépendances et les paramètres populaires. Ne comptez pas sur ~/tao-core à l'exécution ; les mainteneurs régénèrent les schémas/modèles avant de packager la banque de skills.

Politique d'action Train

Ce modèle est activé pour AutoML au niveau du modèle. Avant de traiter toute demande de stage train, lisez references/skill_info.yaml et résolvez l'override d'exécution à partir soit d'une valeur automl_policy explicite, soit de la demande de workflow de l'utilisateur. Traitez les expressions comme « turn off AutoML », « disable AutoML », « no HPO », ou « plain training » comme automl_policy: off pour cette exécution uniquement ; sinon, utilisez par défaut auto. Quand automl_policy: auto, automl_enabled: true, et que schemas/train.schema.json et references/spec_template_train.yaml sont fournis, routez l'action train par défaut via tao-skill-bank:tao-run-automl avec le skill_dir de ce modèle. Préservez les overrides de workflow/application pour les datasets, specs, répertoires de sortie, paramètres GPU/plateforme, checkpoints parents, et automl_policy. Utilisez l'entraînement direct du modèle uniquement quand automl_policy: off ou le schéma train/modèle fourni manque ; dans le cas du schéma manquant, signalez que AutoML est activé mais non exécutable pour ce modèle jusqu'à ce que les schémas soient générés.

Les actions non-train telles que evaluate, inference, export, et les flux de déploiement restent dans cette skill de modèle. L'override automl_policy par exécution ne change pas les métadonnées du modèle.

Exigences d'entraînement

  • Type de dataset : ocrnet
  • Formats : default
  • Métrique de suivi : val_acc

Exigences de Dataset par Action

Action Clé Spec Source Fichiers Liste ?
dataset_convert dataset_convert.input_img_dir id Non
dataset_convert dataset_convert.gt_file id Non
evaluate dataset.character_list_file eval_dataset character_list Non
evaluate evaluate.test_dataset_dir eval_dataset results/{dataset_convert_job_id}/dataset_convert/lmdb Non
export dataset.character_list_file eval_dataset character_list Non
gen_trt_engine gen_trt_engine.tensorrt.calibration.cal_image_dir calibration_dataset Oui
inference dataset.character_list_file eval_dataset character_list Non
inference inference.inference_dataset_dir eval_dataset test.tar.gz Non
prune dataset.character_list_file eval_dataset character_list Non
quantize dataset.train_dataset_dir train_datasets results/{dataset_convert_job_id}/dataset_convert/lmdb Oui
quantize dataset.val_dataset_dir eval_dataset results/{dataset_convert_job_id}/dataset_convert/lmdb Non
quantize dataset.character_list_file eval_dataset character_list Non
quantize dataset.quant_calibration_dataset.images_dir train_datasets train.tar.gz Non
retrain dataset.train_dataset_dir train_datasets results/{dataset_convert_job_id}/dataset_convert/lmdb Oui
retrain dataset.val_dataset_dir eval_dataset results/{dataset_convert_job_id}/dataset_convert/lmdb Non
retrain dataset.character_list_file eval_dataset character_list Non
train dataset.train_dataset_dir train_datasets results/{dataset_convert_job_id}/dataset_convert/lmdb Oui
train dataset.val_dataset_dir eval_dataset results/{dataset_convert_job_id}/dataset_convert/lmdb Non
train dataset.character_list_file eval_dataset character_list Non

Overrides Spec typiques

Les overrides de source de données sont obligatoires pour chaque action — l'agent DOIT construire les chemins des sources de données à partir du tableau Exigences de Dataset par Action ci-dessus et les inclure dans spec_overrides.

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

train (sources de données obligatoires) :

{
    "train.num_epochs": 30,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "train.num_gpus": 1,
    "dataset.batch_size": 16,
    "dataset.train_dataset_dir": [f"{S3_TRAIN}/results/{dataset_convert_job_id}/dataset_convert/lmdb"],
    "dataset.val_dataset_dir": f"{S3_EVAL}/results/{dataset_convert_job_id}/dataset_convert/lmdb",
    "dataset.character_list_file": f"{S3_EVAL}/character_list",
}

gen_trt_engine (sources de données obligatoires) :

{
    "gen_trt_engine.tensorrt.data_type": "fp16",
    "gen_trt_engine.tensorrt.calibration.cal_image_dir": [f"{S3_TRAIN}"],
}

evaluate (sources de données obligatoires) :

{
    "dataset.character_list_file": f"{S3_EVAL}/character_list",
    "evaluate.test_dataset_dir": f"{S3_EVAL}/results/{dataset_convert_job_id}/dataset_convert/lmdb",
}

export (sources de données obligatoires) :

{
    "dataset.character_list_file": f"{S3_EVAL}/character_list",
}

inference (sources de données obligatoires) :

{
    "dataset.character_list_file": f"{S3_EVAL}/character_list",
    "inference.inference_dataset_dir": f"{S3_EVAL}/test.tar.gz",
}

prune (sources de données obligatoires) :

{
    "dataset.character_list_file": f"{S3_EVAL}/character_list",
}

quantize (sources de données obligatoires) :

{
    "dataset.train_dataset_dir": [f"{S3_TRAIN}/results/{dataset_convert_job_id}/dataset_convert/lmdb"],
    "dataset.val_dataset_dir": f"{S3_EVAL}/results/{dataset_convert_job_id}/dataset_convert/lmdb",
    "dataset.character_list_file": f"{S3_EVAL}/character_list",
    "dataset.quant_calibration_dataset.images_dir": f"{S3_TRAIN}/train.tar.gz",
}

retrain (sources de données obligatoires) :

{
    "dataset.train_dataset_dir": [f"{S3_TRAIN}/results/{dataset_convert_job_id}/dataset_convert/lmdb"],
    "dataset.val_dataset_dir": f"{S3_EVAL}/results/{dataset_convert_job_id}/dataset_convert/lmdb",
    "dataset.character_list_file": f"{S3_EVAL}/character_list",
}

Dataset d'Évaluation

Optionnel. Données de test fournies sous forme de tarball séparé.

Paramètres importants

  • dataset.character_list_file : Chemin vers la liste de caractères définissant l'ensemble de caractères supportés. Cela détermine la taille du vocabulaire de sortie.
  • model.backbone : ResNet par défaut.
  • model.prediction : Type de décodeur. CTC ou Attn (basé sur l'attention).
  • train.optim.lr : Taux d'apprentissage. Par défaut 1,0 (optimiseur Adadelta). La valeur par défaut élevée est spécifique à Adadelta.
  • dataset.batch_size : Taille de batch par GPU. Par défaut 16.

Multi-GPU / Multi-Node

Méthode de lancement : Gérée par Lightning (processus Python unique, Lightning lance les workers).

Clé Spec Description Par défaut
train.num_gpus Nombre de GPUs 1
train.gpu_ids Indices de périphériques GPU [0]
train.distributed_strategy Nom de la stratégie auto
  • Stratégie : auto pour single-GPU, lit train.distributed_strategy depuis la config en multi-GPU
  • Pas de num_nodes explicite dans le script train — orienté single-node
  • Modèle léger, un seul GPU généralement suffisant

Matériel

Minimum 1 GPU(s), recommandé 1 GPU(s). 8 Go+ de VRAM par GPU. La reconnaissance OCR de texte est légère. Un seul GPU est généralement suffisant.

Motifs d'erreur

dataset_convert requis : Si vous utilisez des images brutes + fichiers gt, lancez d'abord dataset_convert pour générer le format LMDB.

Incompatibilité de liste de caractères : Tous les caractères dans les données d'entraînement doivent être présents dans le fichier character_list.

Inférence des paramètres Spec / Modèle parent

Les mappages d'inférence spécifiques au modèle appartiennent à ce fichier MD, pas à config.json. Les runners générés doivent lire cette section et appliquer les mappages avec les assistants SDK avant create_job(). Cela reflète le flux ancien de microservices infer_params.py.

Mappages d'inférence depuis TAO Core ocrnet.config.json :

Action Champ Spec Fonction d'inférence Sens
dataset_convert results_dir output_dir répertoire des résultats du job actuel
evaluate encryption_key key clé de chiffrement
evaluate evaluate.checkpoint parent_model fichier de modèle déduit du dossier des résultats du job parent
evaluate evaluate.trt_engine parent_model fichier de modèle déduit du dossier des résultats du job parent
evaluate model.pruned_graph_path pruned_model modèle parent élagué
evaluate results_dir output_dir répertoire des résultats du job actuel
export encryption_key key clé de chiffrement
export export.checkpoint parent_model fichier de modèle déduit du dossier des résultats du job parent
export export.onnx_file create_onnx_file chemin ONNX de sortie
export results_dir output_dir répertoire des résultats du job actuel
gen_trt_engine encryption_key key clé de chiffrement
gen_trt_engine gen_trt_engine.onnx_file parent_model fichier de modèle déduit du dossier des résultats du job parent
gen_trt_engine gen_trt_engine.tensorrt.calibration.cal_cache_file create_cal_cache chemin du cache de calibration
gen_trt_engine gen_trt_engine.trt_engine create_engine_file chemin du moteur TensorRT de sortie
gen_trt_engine results_dir output_dir répertoire des résultats du job actuel
inference encryption_key key clé de chiffrement
inference inference.checkpoint parent_model fichier de modèle déduit du dossier des résultats du job parent
inference inference.trt_engine parent_model fichier de modèle déduit du dossier des résultats du job parent
inference model.pruned_graph_path pruned_model modèle parent élagué
inference results_dir output_dir répertoire des résultats du job actuel
prune encryption_key key clé de chiffrement
prune prune.checkpoint parent_model fichier de modèle déduit du dossier des résultats du job parent
prune prune.pruned_file create_pth_file chemin PTH de sortie
prune results_dir output_dir répertoire des résultats du job actuel
quantize encryption_key key clé de chiffrement
quantize quantize.model_path parent_model fichier de modèle déduit du dossier des résultats du job parent
quantize results_dir output_dir répertoire des résultats du job actuel
retrain encryption_key key clé de chiffrement
retrain model.pruned_graph_path parent_model fichier de modèle déduit du dossier des résultats du job parent
retrain results_dir output_dir répertoire des résultats du job actuel
train encryption_key key clé de chiffrement
train results_dir output_dir répertoire des résultats du job actuel
train train.pretrained_model_path ptm_if_no_resume_model PTM quand aucun checkpoint de reprise n'existe
train train.resume_training_checkpoint_path resume_model fichier de modèle déduit du dossier des résultats du job actuel

Pour parent_model ou parent_model_folder, passez l'id du job enfant train/export/AutoML en amont comme parent_job_id. Le SDK liste le dossier des résultats parents, filtre les artefacts de checkpoint, et retourne le fichier ou dossier de modèle sélectionné. N'ajoutez pas ces mappages à config.json et ne patchez pas les scripts de runner générés pour deviner les chemins des checkpoints.

Déploiement

  • tao-deploy-ocrnet — Flux de déploiement OCRNet pour la génération de moteur TensorRT, l'évaluation TensorRT, et l'inférence TensorRT utilisant TAO Deploy.

Skills similaires