Can't load pretrained TrOCR model

@NielsRogge I get this error when I try to load a local TrOCR checkpoint.

>>> processor = TrOCRProcessor.from_pretrained("./checkpoint-2")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/users/gpupro/gpu_tazi/.local/lib/python3.8/site-packages/transformers/processing_utils.py", line 186, in from_pretrained
    args = cls._get_arguments_from_pretrained(pretrained_model_name_or_path, **kwargs)
  File "/usr/users/gpupro/gpu_tazi/.local/lib/python3.8/site-packages/transformers/processing_utils.py", line 230, in _get_arguments_from_pretrained
    args.append(attribute_class.from_pretrained(pretrained_model_name_or_path, **kwargs))
  File "/usr/users/gpupro/gpu_tazi/.local/lib/python3.8/site-packages/transformers/models/auto/tokenization_auto.py", line 544, in from_pretrained
    tokenizer_class_py, tokenizer_class_fast = TOKENIZER_MAPPING[type(config)]
  File "/usr/users/gpupro/gpu_tazi/.local/lib/python3.8/site-packages/transformers/models/auto/auto_factory.py", line 564, in __getitem__
    raise KeyError(key)
KeyError: <class 'transformers.models.vision_encoder_decoder.configuration_vision_encoder_decoder.VisionEncoderDecoderConfig'>

This is the content of my checkpoint folder:

checkpoint-2
  |-trainer_state.json
  |-preprocessor_config.json
  |-training_args.bin
  |-scaler.pt
  |-optimizer.pt
  |-scheduler.pt
  |-pytorch_model.bin
  |-rng_state.pth
  |-config.json

Yet, loading a TrOCR checkpoint from the hub works just fine.

Issue Analytics

State:
Created a year ago
Comments:13 (2 by maintainers)

Top GitHub Comments

2reactions

ydshiehcommented, May 18, 2022

Hi, @CuarteroAlvaro

In you code snippet, there is 1 line missing

processor.save_pretrained('./local_checkpoint')  # <-- this is required

The following will work

from transformers import TrOCRProcessor, VisionEncoderDecoderModel

processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-printed')
model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-printed')
model.save_pretrained("./local_checkpoint")
processor.save_pretrained('./local_checkpoint')  # <-- this is required 
loaded_preprocessor = TrOCRProcessor.from_pretrained('./local_checkpoint') 
loaded_model = VisionEncoderDecoderModel.from_pretrained('./local_checkpoint')
print(loaded_preprocessor, loaded_model)

But I can’t reproduce the error shown in @NouamaneTazi 's original issue.

KeyError: <class 'transformers.models.vision_encoder_decoder.configuration_vision_encoder_decoder.VisionEncoderDecoderConfig'>

2reactions

emigomezcommented, May 18, 2022

The problem was that my model is on the GPU, but my data is on the CPU. So, I need to send my data to GPU changing on the previous code that I have shared:

generated_ids = model.generate(pixel_values.to(device))

And that is working.

Thank you!!

Top Results From Across the Web

TrOCR - Hugging Face

In this paper, we propose an end-to-end text recognition approach with pre-trained image Transformer and text Transformer models, namely TrOCR, which leverages ...

arXiv:2109.10282v5 [cs.CL] 6 Sep 2022

In this paper, we propose an end-to-end text recognition approach with pre-trained image Transformer and text Transformer models, namely TrOCR, ...

TrOCR: Transformer-based Optical Character ... - DeepAI

The TrOCR model is simple but effective, and can be pre-trained with large-scale synthetic data and fine-tuned with human-labeled datasets.

(PDF) TrOCR: Transformer-based Optical Character ...

The TrOCR model is simple but effective, and can be pre-trained with large-scale synthetic data and fine-tuned with human-labeled datasets.

Load a pre-trained model from disk with Huggingface ...

load(resolved_archive_file, map_location="cpu") except Exception: raise OSError( > "Unable to load weights from pytorch checkpoint file. " "If ...