Size mismatch for linear.weight: copying a param with shape torch.Size([114, 256]) from checkpoin

[janrom1]

I am trying to run doctr with polish vocabs. What I did:

1. I have already changed vocabs.py adding:

VOCABS["polish"] = VOCABS["digits"] + "aąbcćdeęfghijklłmnńoópqrsśtuvwxyzźż" + \
    "AĄBCĆDEĘFGHIJKLŁMNŃOÓPQRSŚTUVWXYZŹŻ" + \
    "!\"#$%&'()*+,-./\:;<=>?@[]^_`{|}~ "

2. I have created 1000 images with trdg. After I ran my train command:

python3 references/recognition/train_pytorch.py crnn_vgg16_bn --train_path /home/roman/learn/doctr-tests/recognition_train --val_path /home/roman/learn/doctr-tests/recognition_val/ --epochs 5 --vocab polish --name crnn_polish_roman

with result

Namespace(arch='crnn_vgg16_bn', train_path='/home/roman/learn/doctr-tests/recognition_train', val_path='/home/roman/learn/doctr-tests/recognition_val/', train_samples=1000, val_samples=20, font='FreeMono.ttf,FreeSans.ttf,FreeSerif.ttf', min_chars=1, max_chars=12, name='crnn_polish_roman', epochs=5, batch_size=64, device=None, input_size=32, lr=0.001, weight_decay=0, workers=None, resume=None, vocab='polish', test_only=False, show_samples=False, wb=False, push_to_hub=False, pretrained=True, sched='cosine', amp=False, find_lr=False)
Validation set loaded in 0.001204s (299 samples in 5 batches)
Downloading https://doctr-static.mindee.com/models?id=v0.3.1/crnn_vgg16_bn-9762b0b0.pt&src=0 to /home/roman/.cache/doctr/models/crnn_vgg16_bn-9762b0b0.pt
63287296it [00:05, 11592246.04it/s]                                                                                     
WARNING:root:No accessible GPU, targe device set to CPU.
Train set loaded in 0.003956s (999 samples in 15 batches)
Validation loss decreased inf --> 39.2596: saving state...                                                              
Epoch 1/5 - Validation loss: 39.2596 (Exact: 0.00% | Partial: 0.00%)
Validation loss decreased 39.2596 --> 5.50948: saving state...                                                          
Epoch 2/5 - Validation loss: 5.50948 (Exact: 0.00% | Partial: 0.00%)
Validation loss decreased 5.50948 --> 5.13989: saving state...                                                          
Epoch 3/5 - Validation loss: 5.13989 (Exact: 0.00% | Partial: 0.00%)
Validation loss decreased 5.13989 --> 4.9927: saving state...                                                           
Epoch 4/5 - Validation loss: 4.9927 (Exact: 0.00% | Partial: 0.00%)
Validation loss decreased 4.9927 --> 4.9836: saving state...                                                            
Epoch 5/5 - Validation loss: 4.9836 (Exact: 0.00% | Partial: 0.00%)

3. than I ran recognition process with script below:

import sys, os
os.environ['USE_TORCH'] = '1'

import torch
from doctr.io import DocumentFile
from doctr.models.recognition import crnn_vgg16_bn
from doctr.models.recognition.predictor import RecognitionPredictor
from doctr.models.preprocessor import PreProcessor
from doctr.file_utils import is_tf_available, is_torch_available

print(f"TF available: {is_tf_available()}, PyTorch available: {is_torch_available()}")

doc = DocumentFile.from_images("invoices/0.jpg")
reco_model = crnn_vgg16_bn(pretrained=False, pretrained_backbone=False, vocab="polish")
reco_model.load_state_dict(torch.load('/home/roman/learn/doctr/crnn_polish_roman.pt', map_location='cpu'))
reco_predictor = RecognitionPredictor(
        PreProcessor((32, 128), preserve_aspect_ratio=True, batch_size=32, mean=(0.694, 0.695, 0.693),
                     std=(0.299, 0.296, 0.301)), reco_model)

predictor = reco_predictor(reco_predictor)
result = predictor(doc)
print(result.export())

unfortunate I got error and I have no idea where to look for:

TF available: False, PyTorch available: True

Traceback (most recent call last):
  File "/mnt/backup-480-ssd/learn/doctr-tests/test0.py", line 24, in <module>
    reco_model.load_state_dict(torch.load('/home/roman/learn/doctr/crnn_polish_roman.pt', map_location='cpu'))
  File "/home/roman/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1497, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for CRNN:
	size mismatch for linear.weight: copying a param with shape torch.Size([114, 256]) from checkpoint, the shape in current model is torch.Size([7, 256]).
	size mismatch for linear.bias: copying a param with shape torch.Size([114]) from checkpoint, the shape in current model is torch.Size([7]).

Do you have any idea what should I change? Where too look up?
My env:

Python 3.10.6
pytorch-lightning 1.8.0.post1
torch 1.11.0+cpu
torchaudio 0.11.0+cpu
torchmetrics 0.10.2
torchvision 0.12.0

[felixdittrich92]

from doctr.datasets import VOCABS

crnn_vgg16_bn(..., vocab=VOCABS["polish"])

👍

[felixdittrich92]

Btw. your model hasn't trained well (exact match is zero) you should trigger the run with --pretrained

[janrom1]

I know ;) It was just for quick check. Thanks a lot for hint with VOCABS.