OCRNet#

OCRNet is a model to recognize characters in an image. It supports the following tasks:

dataset_convert
train
evaluate
prune
inference
export
quantize

Each task is explained in detail in the following sections.

Preparing the Dataset#

The train dataset and evaluation dataset for OCRNet is in LMDB format. You can use dataset_convert to convert the original images and labels to LMDB format. The original dataset should be organized in the following structure:

/Dataset
    /images
        0000.jpg
        0001.jpg
        0002.jpg
        ...
    gt_list.txt
    characters_list.txt

The gt_list.txt file contains all the ground truth text for the images, and each image and its corresponding label is specified with one line of text:

jpg abc
jpg defg
jpg zxv
...

There is a characters_list.txt file that contains all the characters found in the dataset. Each character occupies one line.

Creating an Experiment Specification File#

The experiment specification file includes arguments for all the tasks supported by OCRNet (train/evaluate/inference/prune/export). Here is an example specification file used in the OCRNet get_started notebook:

Parameter	Data Type	Default	Description	Supported Values
`model`	dict config	–	The configuration of the model architecture
`dataset`	dict config	–	The configuration of the dataset
`train`	dict config	–	The configuration of the training task
`evaluate`	dict config	–	The configuration of the evaluation task
`inference`	dict config	–	The configuration of the inference task
`encryption_key`	string	None	The encryption key to encrypt and decrypt model files
`results_dir`	string	/results	The directory where experiment results are saved
`prune`	dict config	–	The configuration for the pruning
`export`	dict config	–	The configuration of the export
`dataset_convert`	dict config	–	The configuration for the dataset convert

model#

The model parameter provides options to change the architecture of OCRNet.

model:
  TPS: True
  backbone: ResNet
  feature_channel: 512
  sequence: BiLSTM
  hidden_size: 256
  prediction: CTC
  quantize: False

Parameter	Datatype	Default	Description	Supported Values
`TPS`	Boolean	False	A flag that enables Thin-plate spline interpolation for the OCRNet input	True/False
`num_fiducial`	Unsigned int	20	The number of fiducial points for TPS	>4
`backbone`	String	ResNet	The backbone of the OCRNet model	ResNet, ResNet2X, FAN_tiny_2X
`feature_channel`	Unsigned int	512	The number of channels for the backbone output feature	>0
`sequence`	String	BiLSTM	The sequence module of the OCRNet model	BiLSTM
`hidden_size`	Unsigned int	256	The number of channels for the BiLSTM hidden layer	>0
`prediction`	String	CTC	The method for encoding and decoding the output feature	CTC, Attn
`input_width`	Unsigned int	100	The input image width	>4
`input_height`	Unsigned int	32	The input image height	>32
`input_channel`	Unsigned int	1	The input image channel	1,3
`quantize`	Boolean	False	A flag that enables quantize and dequantize nodes in the OCRNet backbone	True/False

dataset#

The dataset parameter provides options to set the dataset consumed in training and evaluation.

dataset:
  train_dataset_dir: [/data/train/lmdb]
  val_dataset_dir: /data/test/lmdb
  character_list_file: /data/character_list
  max_label_length: 25
  batch_size: 32
  workers: 4
  augmentation:
    keep_aspect_ratio: False
    aug_prob: 0.3
    reverse_color_prob: 0.5
    rotate_prob: 0.5
    max_rotation_degree: 5
    blur_prob: 0.5
    gaussian_radius_list: [1, 2, 3, 4]

Parameter	Datatype	Default	Description	Supported Values
`train_dataset_dir`	List of String	None	A list of absolute paths to the training datasets. Currently, only a list length of 1 is supported.	List of String
`val_dataset_dir`	String	None	The absolute path to the evaluation dataset	dataset absolute path
`character_list_file`	String	None	The absolute path to character list file	absolute file path
`max_label_length`	Unsigned int	25	The maximum length of the ground truth	>0
`batch_size`	Unsigned int	32	The batch size for training	>0
`workers`	Unsigned int	4	The number of workers to parallel preprocess the training data	>=0
`augmentation`	Dict config	–	The augmentation config.	–

augmentation#

The augmentation parameter provides options to set augmentation pipeline during training.

augmentation:
  keep_aspect_ratio: False
  aug_prob: 0.3
  reverse_color_prob: 0.5
  rotate_prob: 0.5
  max_rotation_degree: 5
  blur_prob: 0.5
  gaussian_radius_list: [1, 2, 3, 4]

Parameter	Datatype	Default	Description	Supported Values
`keep_aspect_ratio`	Bool	False	A flag to enable keeping aspect-ratio when resize the image to model input size	False/True
`aug_prob`	Float	0.0	The probability to apply the following augmentation on the input image	[0, 1]
`reverse_color_prob`	Float	0.5	The probability to reverse the color of the input image	[0, 1]
`rotate_prob`	Float	0.5	The probability to random rotate the input image	[0, 1]
`max_rotation_degree`	Float	0.5	The maximum degree the image will be rotated	>=0
`blur_prob`	Float	0.5	The probability to blur the input image	[0, 1]
`gaussian_radius_list`	List of integer	[1, 2, 3, 4]	The list of radius when apply gaussian blur on the image	–

train#

The train parameter provides options to set training hyperparameters.

train:
  seed: 1111
  gpu_ids: [0]
  optim:
    name: "adadelta"
    lr: 1.0
  clip_grad_norm: 5.0
  num_epochs: 10
  checkpoint_interval: 2
  validation_interval: 1

Parameter	Datatype	Default	Description	Supported Values
`num_gpus`	unsigned int	1	The number of GPUs to use for distributed training	>0
`gpu_ids`	List[int]	[0]	The indices of the GPU’s to use for distributed training
`seed`	unsigned int	1234	The random seed for random, numpy, and torch	>0
`num_epochs`	unsigned int	10	The total number of epochs to run the experiment	>0
`checkpoint_interval`	unsigned int	1	The epoch interval at which the checkpoints are saved	>0
`validation_interval`	unsigned int	1	The epoch interval at which the validation is run	>0
`resume_training_checkpoint_path`	string		The intermediate PyTorch Lightning checkpoint to resume training from
`results_dir`	string	/results/train	The directory to save training results
`optim`	Dict config	–	The configuration for the optimizer	–
`clip_grad_norm`	Float	5.0	The threshold value of magnitude of the gradient L2 norm to be clipped	>4
`distributed_strategy`	String	ddp	The distributed strategy for multi-GPU training	ddp
`pretrained_model_path`	String	None	The absolute path to pretrained weights	–
`quantize_model_path`	String	None	The absolute path to pretrained models for quantize-aware-training	–
`model_ema`	Bool	False	Enable model exponential moving average in the training	False/True

optim#

The optim provides the options to set the optimizer for the training.

optim:
  name: "adadelta"
  lr: 1.0

Parameter	Datatype	Default	Description	Supported Values
`name`	String	adadelta	The optimizer type	adadelta, adam
`lr`	Float	1.0	The initial learning rate for the training	>0.0

evaluate#

The evaluate parameter provides options to set evaluation hyperparameters.

evaluate:
  checkpoint: "??"
  test_dataset_dir: "??"
  results_dir: "${results_dir}/evaluate"

Parameter	Datatype	Default	Description	Supported Values
`checkpoint`	String	–	The absolute path to the model checkpoint for evaluation	–
`results_dir`	String	/results/evaluate	The directory to save evaluation results
`num_gpus`	Unsigned int	1	The number of GPUs to use for distributed evaluation	>0
`gpu_ids`	List[int]	[0]	The indices of the GPU’s to use for distributed evaluation
`test_dataset_dir`	String	–	The absolute path to the evaluation LMDB dataset	–
`batch_size`	Unsigned int	1	The evaluation batch size	>0

prune#

The prune parameter provides options to set prune hyperparameters.

gpu_id: 0
checkpoint: "??"
results_dir: "${results_dir}/prune"
prune_setting:
  mode: experimental_hybrid
  amount: 0.4
  granularity: 8
  raw_prune_score: L1

Parameter	Datatype	Default	Description	Supported Values
`checkpoint`	String	–	The absolute path to the model checkpoint for pruning	–
`gpu_id`	Unsigned int	0	The GPU device index	A valid gpu index
`results_dir`	String	–	The absolute path to the pruning log	–
`pruned_file`	String	–	The absolute path for storing the pruned model checkpoint	–
`prune_setting`	Dict config	–	The pruning hyperparameters	–

prune_setting#

The prune_setting parameter contains options for the pruning algorithms:

Parameter	Datatype	Default	Description	Supported Values
`mode`	String	amount	The absolute path to the model checkpoint to be pruned: `amount`: Prune the amount ratio of weights according to the importance `threshold`: Prune weights with importance smaller than the threshold value `experimental_hybrid`: Prune weights using a hybrid of `threshold` and `amount`	amount, threshold, experimental_hybrid
`amount`	Float	–	The amount value for `amount` and `experimental_hybrid` mode	[0, 1]
`threshold`	Float	–	The threshold value for threshold mode	>=0
`granularity`	Unsigned int	8	The granularity of the pruned layer. The number of pruned-layer output channels will be a multiple of the granularity.	>0
`raw_prune_score`	Dict config	L1	The method for computing the importance of weights	L1, L2

inference#

The inference parameter provides options for inference.

inference:
  checkpoint: "??"
  inference_dataset_dir: "??"
  results_dir: "${results_dir}/inference"

Parameter	Datatype	Default	Description	Supported Values
`checkpoint`	String	–	The absolute path to the model checkpoint for inference	–
`results_dir`	String	/results/inference	The directory to save inference results
`num_gpus`	Unsigned int	1	The number of GPUs to use for distributed inference	>0
`gpu_ids`	List[int]	[0]	The indices of the GPU’s to use for distributed inference
`inference_dataset_dir`	String	–	The absolute path to the inference images directory	–
`batch_size`	Unsigned int	1	The inference batch size	>0

export#

The export parameter provides export options.

export:
  gpu_id: 0
  checkpoint: "??"
  results_dir: "${results_dir}/export"

Parameter	Datatype	Default	Description	Supported Values
`checkpoint`	String	–	The absolute path to the model checkpoint for export	–
`gpu_id`	Unsigned int	0	The GPU device index	Valid gpu index
`onnx_file`	String	–	The absolute path to export ONNX file	–
`results_dir`	String	–	The absolute path to the export output	–

dataset_convert#

The dataset_convert parameter provides options to set dataset conversion.

dataset_convert:
  input_img_dir: "??"
  gt_file: "??"
  results_dir: "${results_dir}/convert_dataset"

Parameter	Datatype	Default	Description	Supported Values
`input_img_dir`	String	–	The absolute path to images directory	–
`gt_file`	String	–	The absolute path to the ground truth file	–
`results_dir`	String	–	The absolute path to `dataset_convert` (i.e. the LMDB dataset and log)	–

Evaluating the Model#

Multi-GPU evaluation is currently not supported for OCRNet.

Inference with the Model#

Multi-GPU inference is currently not supported for OCRNet.

Quantization#

OCRNet supports PTQ via TAO Quant using either the torchao (weight-only) or modelopt (static PTQ) backends.

Add a quantize section to your experiment specification (see TAO Quant documentation for schema and backend options).
Use the quantized checkpoint by setting evaluate.is_quantized: true or inference.is_quantized: true and pointing to the artifact saved under results_dir (for example, quantized_model_torchao.pth or quantized_model_modelopt.pth). For ModelOpt artifacts, the model weights are stored under model_state_dict.

Notes#

For modelopt static PTQ, ensure that your dataset configuration provides a representative calibration loader.
For torchao, activation settings in the configuration are ignored.

Calibration Dataset (ModelOpt)#

When you use the modelopt backend (static PTQ), provide a calibration dataset via dataset.quant_calibration_dataset.

Minimal example:

quantize:
  backend: "modelopt"
  mode: "static_ptq"
  algorithm: "minmax"
dataset:
  quant_calibration_dataset:
    images_dir: "/path/to/calib/images"

See also: TAO Quant overview and its Configuration and backend pages.

Deploying to DeepStream#

For DeepStream integration, see Deploy nvOCDR to DeepStream.