How to finetune my trained text detection model for KIE

[sophiadouglas]

Hi,

I have trained a custom detection model from the pretrained fast_base model using the following commands:

python train_tensorflow.py fast_base --train_path=E:\test\train --val_path=E:\test\val --epochs=100 --batch_size=1 --freeze-backbone --early-stop --early-stop-epochs=50 --save-interval-epoch --eval-straight --rotation --pretrained

A desirable result is reached with Recall: 96%, Precision: 92%, Mean IoU: 78%.

Now, I would like to finetune this model to cater for KIE detection with 8 new classes.

I have tried using --pretrained-backbone, --resume, and --freeze-backbone but can't seem to make it work.

May I know how to go about from here? What arguments should I use to finetune my trained detection model for KIE?

Thank you so much.

[felixdittrich92]

Hi @sophiadouglas 👋,

Glad to hear you have fine tuned your custom model. 👏

You follow the dataset format as described here ?
The KIE training requires lots (really lots) of data. How huge is your dataset ?

Maybe also a way which could be interesting: #1669 (comment)

CC @aminemindee

Best regards,
Felix

[sophiadouglas]

Hi @felixdittrich92 😄

Thank you for the reply.

I have followed the dataset format as described.
For testing purposes my training dataset is over 400 images only.

Using the aforementioned training dataset, I have actually managed to train my custom KIE model from the available pretrained detection model with sufficient results (just enough to confirm that the training script is used correctly).

However, the task that I am interested now is training a custom KIE model from my trained detection model.

As an example, let me describe my situation this way:
From the available pretrained model (tensorflow fast_base), I have finetuned the model to detect text from a student id card dataset. Now, from this custom detection model, I would like to retrain it to cater for a KIE task which detects multiple classes like name and student id etc.

How should I finetune my already trained text detection model (which was trained on one class) to a custom KIE model (to be trained on multiple classes)?

I am guessing I need to use --pretrained-backbone (path to the backbone weights) but I am not sure how to obtain these weights from my trained detection model as the output is its checkpoint.

I hope this makes sense.

Any help is much appreciated.

Best regards,
Sophia

PS Thank you very much for the referenced comment #1669.
I have also tried the suggestion with GliNER before this.
However, it does not work well for our custom task that's why we opted to train our own model.

[felixT2K]

Hi @sophiadouglas 👋🏼 ,

What you need is --resume which loads the whole model and "hopefully" resets the head to avoid a shape mismatch while init the weights.
Otherwise try to set skip_mismatch=True to model.load_weights.

I remember we have had trouble with tf about reseting the head not only for KIE also for fine tuning the recognition models on a different vocab like the pretrained one.

Best,
Felix

[sophiadouglas]

Hi @felixT2K 👋

Thank you for the reply.

I have tried --resume and setting skip_mismatch=True to model.load_weights and it gave me this error:
ValueError: When calling model.load_weights, skip_mismatch can only be set to True when by_name is True.

I later added by_name=True then it gave me this:
ValueError: Received incompatible tensor with shape (1, 1, 128, 1) when attempting to restore variable with shape (1, 1, 128, 8) and name head/layer_with_weights-1/kernel/.ATTRIBUTES/VARIABLE_VALUE.

I remember we have had trouble with tf about reseting the head not only for KIE also for fine tuning the recognition models on a different vocab like the pretrained one.

Do you know if this has been resolved in tf? Would it work in pytorch instead?

Thank you so much.

Best regards,
Sophia

[felixdittrich92]

Hi,

in pytorch it works yes (ref.:

doctr/doctr/models/detection/fast/pytorch.py

Line 348 in 9045dcf

# The number of class_names is not the same as the number of classes in the pretrained model =>

) it will reset the head

With tf we can't do a lot until we switch the saving format to .h5 .. :/ :)

[sophiadouglas]

Hi @felixdittrich92 👋

I see, alright I will try it out with pytorch.

Thank you very much for your help.

Best regards,
Sophia