Ideal Podium

Optimizing Parseq OCR: Adding Quantization Support for Enhanced Efficiency

Vikas Kumar Ojha

MachineLearning@Avalara Ex-Samsung

In an era driven by information and data, the ability to extract text from images is a technology that has reshaped the way we interact with both physical and digital content. Optical Character Recognition, commonly known as OCR, is the ingenious technology that empowers machines to decipher and convert printed or handwritten text from various sources into editable and searchable digital text. There are numerous applications of OCR in the industry. In this blog, we will learn how to train and quantize one of the best OCR engines available in the open source.

----

Introduction

Recently I was exploring the open source options of OCR for one of my personal projects and through my references I came across a model called Parseq. That is based on vision transformers and uses language modeling to recognize the text which is difficult to be recognized with the image features alone. The model has encoder-decoder-based architecture with multi-headed attention, which helps achieve good accuracy and low compute requirements. However for deploying on mobile or edge devices the more we can optimize the better it is as these devices generally have pretty low processing capabilities.

Recently I experimented with Parseq-tiny and found its result was pretty good on my dataset. Hence, I began to find a way to quantize the model in order to reduce its size as well its compute requirement but I couldn’t find any leads from the official implementation. Hence I decided to implement the same. I did certain experiments and found the correct recipe to optimize the model. After doing a comparison I found the implementation to be working well hence I integrated the same with a fork of the repo. We would use the same in the subsequent sections and see how you can use the same to optimize the model for your use cases.

----

Experimentation and Analysis

Official documentation from Pytorch as well as certain research recommends using dynamic quantization for transformer models as it determines the scale and zero values for quantization of weights from the inference data itself allowing the weights and activations to be quantized on the fly allowing to preserve accuracy. On the other hand in static quantization generally, there is a loss of accuracy especially if the model is not trained in QAT fashion because of loss in information due to typecasting.

Hence I decided to perform dynamic quantization without QAT (Quantization Aware Training) as QAT was not available in the default code. Again I added QAT support to the training code and retrained the model. For the first time, I trained the model for around 25 epochs and epoch 22 had the best accuracy. With QAT added I tried training it for a little longer around 30 epochs and in that case, the best epoch was around 23. Below is the table summarizing the experimentation.

In the case of dynamic quantization, the accuracy drop is minimal, and hence the difference in accuracy drop/error rate increase is not huge but still it can be seen in the above table that in the case of quantization without QAT the Character Error Rate rose by 0.00208% while it raised by 0.00131 in case of QAT trained model. The Accuracy which here means having an exact word-to-word match reduced in the first case by 0.2015 while it decreased by 0.1484 for the QAT model which proves QAT has helped in this case.

----

Training and Optimizing the Model on Your Own Dataset

Having seen the comparison between the quantized and non-quantized parseq model. Let’s use my repository to train a parseq-tiny model and perform dynamic quantization on the same.

Downloading the dataset.

For the sake of demonstration, I will use the IAM words dataset for training our OCR model. It is a dataset that consists of handwritten words. So let's go ahead and download and extract the same.

          !wget -q https://github.com/sayakpaul/Handwriting-Recognizer-in-Keras/releases/download/v1.0.0/IAM_Words.zip
!unzip /content/IAM_Words.zip

----

Converting the dataset into LMDB format

By default, the parseq training implementation uses the LMDB data format for training the model. Hence we need to convert the download dataset into LMBD format. Hence let’s go ahead and do it.

We would be using a great repo by clovaai to convert our dataset into the LMDB format. The repo requires the dataset to be prepared in a certain fixed format in order to perform the conversion hence we need to convert our dataset in that format before we convert to the LMDB format which is demonstrated below.

The images should be kept in a folder inside a folder and in the parent folder there should be a ground truth text file containing the labels of each image as well as the path of the images as shown above. Once that is done you can use the create_lmdb_dataset.py script in the repo that converts our dataset into LMDB format. For a detailed understanding of the process please see my notebook provided at the end of this blog. I skipped that because that would make the blog too long. Also, the steps might vary for different datasets.

----

Training parseq tiny in a quantization aware setting

Here we would use my version of the parseq and perform Quantization Aware of parseq tiny.

Cloning my repo and installing the dependencies required for training

          !git clone https://github.com/VikasOjha666/parseq.git
%cd parseq
!pip install -r requirements.txt

%cd parseq

----

Prepare the dataset for training

Assuming the data exists in folder iam_dataset. We need to arrange the data in a certain format required by the parseq training code hence let's go ahead and do it.

          !mkdir train
!mkdir train/real
!mkdir train/real/iam/

!mkdir val
!mkdir val/real
!mkdir val/real/iam/


!cp -r /content/iam_dataset/train/* train/real/iam/
!cp -r /content/iam_dataset/val/* val/real/iam/

----

Editing the configuration file to train parseq tiny

          %%writefile ./configs/model/parseq.yaml
name: parseq
_target_: strhub.models.parseq.system.PARSeq

# Data
patch_size: [ 4, 8 ]  # [ height, width ]

# Architecture
embed_dim: 192
enc_num_heads: 3
enc_mlp_ratio: 4
enc_depth: 12
dec_num_heads: 6
dec_mlp_ratio: 4
dec_depth: 1

# Training
lr: 7e-4
perm_num: 6
perm_forward: true
perm_mirrored: true
dropout: 0.1

# Decoding mode (test)
decode_ar: true
refine_iters: 1

Here we edit the configuration file of parseq setting the parameters like embed_dim, enc_num_heads which determines the structure of the neural network.

----

Run the training

          !python ./train.py dataset=real data.batch_size=64 model.batch_size=64 trainer.gpus=1 data.root_dir=./ trainer.val_check_interval=1000 data.remove_whitespace=false pretrained=parseq-tiny trainer.max_epochs=5

Here I am training the model for just 5 epochs you can set more if you want better results.

----

Quantizing the model

          !python quantize_model.py --source_path /content/parseq/outputs/parseq/2023-08-30_12-20-56/checkpoints/epoch=4-step=5220-val_accuracy=89.4386-val_NED=96.7970.ckpt --target_path /content/  --optimize_for_mobile False

Quantizing the model is as simple as running this script which I have created for performing dynamic quantization on the QAT trained model. Below is the explanation of the parameters of the script.

source_path: This is the path to the checkpoint of the trained model.

target_path: This is the directory in which the optimized model should be saved.

optimize_for_mobile: This is by default True and setting this optimizes the model for mobile inferences you can set it to false for inference on x86 platforms.

----

References and Further Readings

Scene Text Recognition with Permuted Autoregressive Sequence Models

Context-aware STR methods typically use internal autoregressive (AR) language models (LM). Inherent limitations of AR models motivated two-stage methods which employ an external LM. The conditional independence of the external LM on the input image may cause it to erroneously rectify correct predictions, leading to significant inefficiencies. Our method, PARSeq, learns an ensemble of internal AR LMs with shared weights using Permutation Language Modeling. It unifies context-free non-AR and context-aware AR inference, and iterative refinement using bidirectional context. Using synthetic training data, PARSeq achieves state-of-the-art (SOTA) results in STR benchmarks (91.9% accuracy) and more challenging datasets. It establishes new SOTA results (96.0% accuracy) when trained on real data. PARSeq is optimal on accuracy vs parameter count, FLOPS, and latency because of its simple, unified structure and parallel token processing. Due to its extensive use of attention, it is robust on arbitrarily-oriented text which is common in real-world images. Code, pretrained weights, and data are available at: https://github.com/baudm/parseq.

arxiv.org