Load a custom trained KIE detection model Issue

[HuSSonO3]

Hayoo again,

I have an issue where after training the dataset with this command

python train_pytorch.py --train_path ./train --val_path ./valid db_resnet50 -j 0

-j 0 fixed another issue i had where it produces this error

Validation set loaded in 0.009015s (6 samples in 3 batches)
WARNING:root:No accessible GPU, target device set to CPU.
Train set loaded in 0.01067s (15 samples in 7 batches)
                                                                                                                                          Traceback (most recent call last):                                                                                   | 0/7 [00:00<?, ?it/s] 
  File "E:\Work\Receipt Center\OCR\ocr-scripts\train_pytorch.py", line 481, in <module>
    main(args)
  File "E:\Work\Receipt Center\OCR\ocr-scripts\train_pytorch.py", line 388, in main
    fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, amp=args.amp)
  File "E:\Work\Receipt Center\OCR\ocr-scripts\train_pytorch.py", line 108, in fit_one_epoch
    pbar = tqdm(train_loader, position=1)
  File "E:\Work\Receipt Center\OCR\ocr-scripts\OcrEnv\lib\site-packages\tqdm\asyncio.py", line 33, in __init__
    self.iterable_iterator = iter(iterable)
  File "E:\Work\Receipt Center\OCR\ocr-scripts\OcrEnv\lib\site-packages\torch\utils\data\dataloader.py", line 484, in __iter__
    return self._get_iterator()
  File "E:\Work\Receipt Center\OCR\ocr-scripts\OcrEnv\lib\site-packages\torch\utils\data\dataloader.py", line 415, in _get_iterator        
    return _MultiProcessingDataLoaderIter(self)
  File "E:\Work\Receipt Center\OCR\ocr-scripts\OcrEnv\lib\site-packages\torch\utils\data\dataloader.py", line 1138, in __init__
    w.start()
  File "E:\Python\lib\multiprocessing\process.py", line 121, in start
    self._popen = self._Popen(self)
  File "E:\Python\lib\multiprocessing\context.py", line 224, in _Popen
    return _default_context.get_context().Process._Popen(process_obj)
  File "E:\Python\lib\multiprocessing\context.py", line 327, in _Popen
    return Popen(process_obj)
  File "E:\Python\lib\multiprocessing\popen_spawn_win32.py", line 93, in __init__
    reduction.dump(process_obj, to_child)
  File "E:\Python\lib\multiprocessing\reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
AttributeError: Can't pickle local object 'main.<locals>.<lambda>'
  0%|

which also for some reason even that it works it stops producing Recall, Precision and Mean IoU after 10 ephocs

Anyway the issue is after i load the model in here:

import torch
from doctr.models import kie_predictor, db_resnet50
from doctr.io import DocumentFile

det_model = db_resnet50(pretrained=False, pretrained_backbone=False, class_names=['total', 'date', 'item'])
det_params = torch.load('db_resnet50_20241126-161226.pt', map_location="cpu")
det_model.load_state_dict(det_params)

model = kie_predictor(det_arch=det_model, reco_arch='crnn_vgg16_bn', pretrained=True)
doc = DocumentFile.from_images("images/Test/test17.jpg")
result = model(doc)
result.show()

which in turn produces this error:

Traceback (most recent call last):
  File "E:\Work\Receipt Center\OCR\ocr-scripts\docr2.py", line 14, in <module>
    result = model(doc)
  File "E:\Work\Receipt Center\OCR\ocr-scripts\OcrEnv\lib\site-packages\torch\nn\modules\module.py", line 1736, in _wrapped_call_impl      
    return self._call_impl(*args, **kwargs)
  File "E:\Work\Receipt Center\OCR\ocr-scripts\OcrEnv\lib\site-packages\torch\nn\modules\module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "E:\Work\Receipt Center\OCR\ocr-scripts\OcrEnv\lib\site-packages\torch\utils\_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "E:\Work\Receipt Center\OCR\ocr-scripts\OcrEnv\lib\site-packages\doctr\models\kie_predictor\pytorch.py", line 108, in forward       
    dict_loc_preds: Dict[str, List[np.ndarray]] = invert_data_structure(loc_preds)  # type: ignore[assignment]
  File "E:\Work\Receipt Center\OCR\ocr-scripts\OcrEnv\lib\site-packages\doctr\models\_utils.py", line 198, in invert_data_structure        
    return {k: [dic[k] for dic in x] for k in x[0]}
  File "E:\Work\Receipt Center\OCR\ocr-scripts\OcrEnv\lib\site-packages\doctr\models\_utils.py", line 198, in <dictcomp>
    return {k: [dic[k] for dic in x] for k in x[0]}
  File "E:\Work\Receipt Center\OCR\ocr-scripts\OcrEnv\lib\site-packages\doctr\models\_utils.py", line 198, in <listcomp>
    return {k: [dic[k] for dic in x] for k in x[0]}
KeyError: 'date'

This is a snippet of my labels.json since am assuming it has some relation to it.

"test18.jpg": {
        "img_dimensions": [300, 400],
        "img_hash": "8423ee5a47c18930a11690fc57f2de3045f3db78b2c65a6daf5d5ea50ec35c02",
        "polygons": {
            "date": [[[151,303], [205,300], [205,322], [148,317]]],
            "total": [[[65,230], [272,234], [268,251], [63,244]]],  
            "item": [[[54,177], [264,180], [263,197], [55,190]],
            [[53,222], [267,223], [265,207], [51,206]],
            [[52,205], [267,211], [265,197], [54,190]]]
        }
    }

Thank you in advance for your support!

[felixdittrich92]

Hi @HuSSonO3 👋,

1. It's a known issue with PyTorch dataloaders / python multiproc on Windows machines that it makes sometimes trouble with multiple workers (-j), unfortunately we can't do anything here.
2. It stopped after 10 epochs because you haven't provided any value in the train command and our default value is 10 here you can specify it with --epochs 20 for example.
3. Dataset sample looks correct and how it's loaded also ..

I think the issue is that your validation dataset doesn't contains all labels which are in the training set (it should for validation) because we init the model for training with all unique class names from the validation dataset:

doctr/references/detection/train_pytorch.py

Line 222 in 72bed41

class_names=val_set.class_names,

One general note:
Especially for fine tuning on low data i can't suggest to train / fine tune a KIE model with docTR because it uses only visual features instead of multimodal (vis + text) - Personally i would go with docTR as OCR Engine to extract all text and fine tune a LiLT (https://huggingface.co/docs/transformers/model_doc/lilt) model for example which should work better with such an low amount of data and then combine both docTR for OCR and LiLT for KIE extraction.

Best,
Felix