Problem loading trained pt file for ocr recognition.

[neojg]

I have done an experiment training the vitstr_base model for "polish" vocab.

The training has been performed using the following command:

python references/recognition/train.py vitstr_base --pretrained --vocab polish  --output_dir ./train_test --name FirstTrain --font arial.ttf,arialbd.ttf,...,verdanaz.ttf --train-samples 20000 --val-samples 400 --max-chars 18 --epochs 50 --batch_size 192

I copied the checkpoint file for testing.

Now I just like to use it for ocr for test purposes:

reco_model = vitstr_base(pretrained=False, vocab="polish")
reco_model.from_pretrained(
    f"{model_root_folder}FirstTrain.vl.0.0515268.pt", 
    map_location="cpu", 
    vocab="polish", 
)

I'm getting error message:

Message=Error(s) in loading state_dict for ViTSTR:
	size mismatch for head.weight: copying a param with shape torch.Size([119, 768]) from checkpoint, the shape in current model is torch.Size([7, 768]).
	size mismatch for head.bias: copying a param with shape torch.Size([119]) from checkpoint, the shape in current model is torch.Size([7]).

Please advise what I'm doing wrong.

[neojg]

As expected the solution is trivial.

The training script expects the name of the vocabulary e.g. "polish"

The parameter vocab= expects not the name of the language but the actual string containing vocabulary. In my case VOCAB["polish"].

So the working code looks like:

reco_model = vitstr_base(pretrained=False, vocab="polish")
reco_model.from_pretrained(
    f"{model_root_folder}FirstTrain.vl.0.0515268.pt", 
    map_location="cpu", 
    vocab=VOCAB["polish"], 
)