AttributeError: 'DBNet' object has no attribute 'load_state_dict'

[shavejshaikh]

Unable to load a pytorch train model in doctr package getting an error
model = ocr_predictor(det_arch='db_resnet50', reco_arch='sar_resnet31',pretrained=True) model.det_predictor.model.load_state_dict(det_model_path) model.det_predictor.model.postprocessor.unclip_ratio = 2 model.reco_predictor.model.load_weights(rec_model_path)

[felixdittrich92]

See #940 (same way for detection model) If you use tensorflow you need to use load_weights instead :)

[felixdittrich92]

Please note some architectures does not have pretrained checkpoints currently (for example sar_resnet31) we are on it :)

[shavejshaikh]

Hi i have tried #940 but still getting same error

 model.load_state_dict(det_model_path,map_location="cpu")
AttributeError: 'DBNet' object has no attribute 'load_state_dict'

[felixdittrich92]

Could you please check that you have only installed pytorch and not also tensorflow ?

[shavejshaikh]

Yes i have pytorch in requirements.txt also tried running with docker get this error

RuntimeError: Error(s) in loading state_dict for DBNet:

[felixdittrich92]

Ah i see the problem:

model = db_resnet50(pretrained=False)
model.load_state_dict(path, mapping='cpu')
predictor = ocr_predictor(et_arch=model, pretrained=True)

that's the way to load a custom checkpoint :)

[shavejshaikh]

I have tried this code but still getting same error

AttributeError: 'DBNet' object has no attribute 'load_state_dict'

Code

model = db_resnet50(pretrained=False)
model.load_state_dict(torch.load(det_model_path, map_location= "cpu"))
predictor = ocr_predictor(det_arch=model, pretrained=True)

[felixdittrich92]

You have only pytorch installed and not tensorflow correct? Please try to run your code with a USE_TORCH=1 python3 Yourskript.py

[shavejshaikh]

Yes i have install pytorch version of doctr

import torch
import os
from doctr.models import ocr_predictor, db_resnet50
os.environ['USE_TORCH'] = "1"

class DoctrModel:
    def __init__(self, det_model_path=None, rec_model_path=None):
        model = db_resnet50(pretrained=False)
        model.load_state_dict(torch.load(
            det_model_path, map_location="cpu"))
        self.predictor = ocr_predictor(det_arch=model, pretrained=True)

    def __call__(self):
        return self.predictor

Error getting after running it

raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for DBNet:

[shavejshaikh]

Requirements.txt

# Automatically generated by https://github.com/damnever/pigar.

fastapi==0.97.0
pdf2image==1.16.0
Pillow==9.5.0
python-doctr==0.6.0
torch==2.0.1
transformers==4.30.2
uvicorn==0.22.0

[shavejshaikh]

Hi @felixdittrich92 any idea ?

[felixT2K]

Not really :/ Could you run this script please and post the output ? The provided scipt works for me without problems
https://github.com/mindee/doctr/blob/main/scripts/collect_env.py

[shavejshaikh]

Hi, I found a bug when I'm running it as the python script is working fine. But when I add this code in Fastapi it gives me RUNTIME ERROR

[shavejshaikh]

@felixdittrich92 @felixT2K Environment Info

Collecting environment information...


DocTR version: v0.6.0
TensorFlow version: 2.12.0
PyTorch version: 2.0.1+cu117 (torchvision 0.15.2+cu117)
OpenCV version: 4.7.0
OS: Ubuntu 20.04.5 LTS
Python version: 3.9.5
Is CUDA available (TensorFlow): No
Is CUDA available (PyTorch): No
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA

[felixdittrich92]

Hey @shavejshaikh 👋,
As you can see you have installed both... Tensorflow and Pytorch... Doctr will look first if tensorflow is available (that's the reason why load_state_dict does not work).
Especially for production systems i suggest the following:

1. Install your prefered backend: pytorch or Tensorflow
2. Install doctr pip3 install python-doctr

[shavejshaikh]

I found a solution its path issue I'm passing the recognition value to the detection arch i.e. causing mismatch error. Thanks to @felixdittrich92 and @felixT2K for helping me.