flyingbertman commited on
Commit
24e5b2f
Β·
verified Β·
1 Parent(s): 1869e96

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +123 -0
README.md CHANGED
@@ -1,3 +1,126 @@
1
  ---
2
  license: mit
 
 
 
 
 
 
 
 
 
 
 
 
3
  ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
  license: mit
3
+ library_name: onnx
4
+ tags:
5
+ - ocr
6
+ - text-recognition
7
+ - trocr
8
+ - documents
9
+ - transformer
10
+ - onnx
11
+ base_model: microsoft/trocr-base-printed
12
+ pipeline_tag: image-to-text
13
+ language:
14
+ - en
15
  ---
16
+
17
+ # TrOCR Base Printed (ONNX, fp32 + fp16 bundle)
18
+
19
+ ONNX exports of [microsoft/trocr-base-printed](https://huggingface.co/microsoft/trocr-base-printed) β€” Microsoft's Transformer-based OCR for printed text. ViT image encoder + GPT-style autoregressive text decoder, trained on SROIE printed-text crops. Recognizes the text in a cropped image of a single line/word.
20
+
21
+ This repo bundles **both fp32 and fp16 precisions** in one download β€” distribution symmetry, shared tokenizer + config files. Pick a precision via the `.onnx` filename in the `onnx/` subdir.
22
+
23
+ Re-exported from upstream PyTorch weights via a two-step pipeline. Provenance trail: Li et al. β†’ microsoft/trocr-base-printed β†’ `optimum-cli export onnx --task image-to-text-with-past` (fp32 stage) β†’ `onnxconverter_common.float16.convert_float_to_float16(..., keep_io_types=True)` (fp16 cast on the fp32 graph) β†’ these files.
24
+
25
+ Toolchain: `torch 2.4.x` (CUDA 12.4), `transformers 4.45.2`, `optimum[onnxruntime] 1.24.0`, `onnxconverter-common>=1.14`. Full conversion script: [`scripts/export-trocr-base-printed-fp16.ps1`](https://github.com/HeliosophLLC/DatumIngest/blob/main/scripts/export-trocr-base-printed-fp16.ps1) in the DatumIngest repo (despite the `-fp16` suffix on the script name, it produces both precisions in one run).
26
+
27
+ Why the two-step pipeline instead of `optimum-cli ... --dtype fp16` directly: the CUDA path requires a CUDA-enabled torch build in the venv (the other export scripts don't install one), and optimum-cli's fp16 merged-decoder export ships with an If-subgraph wiring bug on this architecture. The onnxconverter-common pass operates on the already-traced fp32 graph in place and sidesteps both issues. `keep_io_types=True` means inputs and outputs stay fp32 at the wire boundary β€” only internal weights + activations run in half precision β€” so runtime code feeds the same input tensors regardless of which `.onnx` file it loads.
28
+
29
+ Credit: Minghao Li, Tengchao Lv, Lei Cui, Yijuan Lu, Dinei Florencio, Cha Zhang, Zhoujun Li, Furu Wei (Microsoft Research). Paper: *"TrOCR: Transformer-based Optical Character Recognition with Pre-trained Models"*, 2021.
30
+
31
+ ## What this repo contains
32
+
33
+ TrOCR is encoder-decoder, so the export splits into multiple files. **All shared (root) files must be present** along with the precision-specific `.onnx` files you choose to load.
34
+
35
+ ### `onnx/` subdir β€” precision-specific files
36
+
37
+ | File | Variant | Size | Role |
38
+ |---|---|---|---|
39
+ | `encoder_model.onnx` | fp32 | ~700 MB | ViT image encoder |
40
+ | `encoder_model_fp16.onnx` | fp16 | ~350 MB | Half-precision ViT image encoder |
41
+ | `decoder_model_merged.onnx` | fp32 | ~700 MB | Text decoder with KV cache merged into one graph |
42
+ | `decoder_model_merged_fp16.onnx` | fp16 | ~350 MB | Half-precision text decoder |
43
+
44
+ The non-merged `decoder_model.onnx` is deliberately omitted β€” the merged form supersedes it for runtime use; keeping both would double the repo size for no benefit.
45
+
46
+ ### Root β€” shared tokenizer + config files
47
+
48
+ | File | Role |
49
+ |---|---|
50
+ | `config.json` | Model architecture config |
51
+ | `generation_config.json` | Decoder generation defaults (max_length, EOS token, etc.) |
52
+ | `preprocessor_config.json` | Image preprocessing β€” `TrOCRProcessor` settings (resize, normalize) |
53
+ | `tokenizer.json` + `vocab.json` + `merges.txt` | BPE tokenizer files |
54
+ | `tokenizer_config.json` + `special_tokens_map.json` | Tokenizer metadata |
55
+
56
+ ## Input / output (both variants)
57
+
58
+ | Stage | Input | Output |
59
+ |---|---|---|
60
+ | Encoder | `pixel_values` β€” NCHW **float32** (yes, even for the fp16 variant β€” IO types are kept fp32), preprocessed RGB image (typically 384Γ—384) | `last_hidden_state` β€” encoder features (fp32 at the boundary; fp16 internally for the half-precision variant) |
61
+ | Decoder | `input_ids` (token sequence so far) + `encoder_hidden_states` + KV cache from prior step | next-token logits + updated KV cache |
62
+
63
+ The fp16 variant's `keep_io_types=True` setting means runtime code is **identical** between fp32 and fp16 β€” you don't have to cast inputs to `np.float16`. Only the on-disk weights and the internal compute differ.
64
+
65
+ ## How to use
66
+
67
+ Greedy decoding orchestrated outside the ONNX graph β€” same encoder-decoder shape as Whisper, T5, BART, and friends:
68
+
69
+ ```python
70
+ import onnxruntime as ort
71
+ import numpy as np
72
+ from PIL import Image
73
+ from transformers import TrOCRProcessor
74
+
75
+ # Pick a precision β€” same runtime code either way thanks to keep_io_types.
76
+ PRECISION_SUFFIX = "" # "" for fp32, "_fp16" for fp16
77
+
78
+ proc = TrOCRProcessor.from_pretrained(".")
79
+ encoder = ort.InferenceSession(f"onnx/encoder_model{PRECISION_SUFFIX}.onnx")
80
+ decoder = ort.InferenceSession(f"onnx/decoder_model_merged{PRECISION_SUFFIX}.onnx")
81
+
82
+ img = Image.open("cropped_text_line.jpg").convert("RGB")
83
+ pixel_values = proc(images=img, return_tensors="np").pixel_values # float32 in both cases
84
+ encoder_hidden = encoder.run(None, {"pixel_values": pixel_values})[0]
85
+
86
+ BOS = proc.tokenizer.cls_token_id
87
+ EOS = proc.tokenizer.eos_token_id
88
+ input_ids = np.array([[BOS]], dtype=np.int64)
89
+ generated, past_kv = [], None
90
+
91
+ for _ in range(64):
92
+ decoder_inputs = {"input_ids": input_ids, "encoder_hidden_states": encoder_hidden}
93
+ if past_kv is not None:
94
+ decoder_inputs.update(past_kv_to_inputs(past_kv))
95
+ outputs = decoder.run(None, decoder_inputs)
96
+ next_token = outputs[0][:, -1, :].argmax(-1)
97
+ if next_token.item() == EOS: break
98
+ generated.append(next_token.item())
99
+ input_ids = next_token.reshape(1, 1)
100
+ past_kv = outputs_to_past_kv(outputs[1:])
101
+
102
+ text = proc.tokenizer.decode(generated, skip_special_tokens=True)
103
+ ```
104
+
105
+ The exact past-KV input/output names are Optimum-version-specific; inspect with Netron once after export to pin them down.
106
+
107
+ ## Which precision should I use?
108
+
109
+ - **fp32** β€” full precision, identical numerics to upstream PyTorch reference. Default for accuracy-sensitive scientific work, OCR-accuracy benchmarks.
110
+ - **fp16** β€” ~half the disk footprint (~700 MB vs ~1.4 GB total) and half the model-load memory. On GPU / NPU with native fp16: modest speedup (typically 1.5-2Γ— over fp32 on consumer GPUs). On CPU runtimes that upcast fp16 β†’ fp32 internally, runtime speed is identical to fp32 but you save the memory.
111
+
112
+ The `keep_io_types=True` setting means switching between them is a single file-path change β€” no code changes needed.
113
+
114
+ ## Related variants (not in this repo)
115
+
116
+ Microsoft publishes a small variant family of TrOCR β€” same architecture, different size / training corpus:
117
+
118
+ - `microsoft/trocr-small-printed` β€” smaller (~5Γ— less), less accurate but faster.
119
+ - `microsoft/trocr-large-printed` β€” bigger, better quality.
120
+ - `microsoft/trocr-base-handwritten` β€” same size as this, trained on IAM handwritten dataset instead of SROIE.
121
+
122
+ All MIT-licensed, all re-exportable via the same script with a swapped `--model` arg.
123
+
124
+ ## License
125
+
126
+ **MIT** β€” same as upstream [microsoft/trocr-base-printed](https://huggingface.co/microsoft/trocr-base-printed). `LICENSE` file included. Optimum ONNX export + fp16 conversion are numerical transformations only β€” no relicensing implication.