Exporting to ONNX
Relevant source files
The following files were used as context for generating this wiki page:
The purpose of exporting to ONNX (Open Neural Network Exchange) in this project is to convert a trained PyTorch OCR model into a standardized format that can be used for efficient inference without requiring a full PyTorch environment. This process bridges the gap between the training phase, which utilizes heavy dependencies like torch and torchvision, and the production/use phase, which relies on onnxruntime for prediction.
By exporting to ONNX, the project enables the OcrModel to be bundled as a static file (model.onnx) and distributed via standard Python packaging tools. This allows end-users to perform CAPTCHA recognition using only lightweight libraries such as numpy and Pillow.
Sources: README.md:1-47, setup.py:20-25, train_model/train_requirements.txt:1-7
Export Process Logic
The export process occurs immediately after the training loop in the training script. It involves transitioning the model from its training state to a format suitable for serialized deployment.
Model Preparation and Trace
Before exporting, the model architecture defined in core.py (a modified ResNet18) is instantiated and trained. To perform the export, PyTorch requires a dummy input tensor that matches the expected input shape of the model to trace the execution graph.
- Input Shape: (1, 1, 50, 160) representing (Batch Size, Channels, Height, Width).
- Export Function:
torch.onnx.export.
The flow shows the transition from a live PyTorch object to a serialized ONNX file.
Sources: train_model/train.py:56-57, train_model/core.py:68-75
Export Configuration
The export command specifies metadata for the resulting ONNX graph, ensuring the input and output nodes are easily identifiable during inference.
| Parameter | Value | Description |
|---|---|---|
model | OcrModel | The trained PyTorch module instance. |
args | (torch_input,) | A tuple containing the dummy input tensor. |
f | "modal.onnx" | The destination filename (noted as modal.onnx in train script, though used as model.onnx elsewhere). |
input_names | ['input'] | Label for the input node. |
output_names | ['output'] | Label for the output node. |
Sources: train_model/train.py:57
Model Architecture for ONNX
The architecture is designed to be compatible with ONNX conversion by using standard layers and explicit reshaping in the forward pass.
Core Components
The model uses a ResNet18 backbone with a custom input layer to handle single-channel (grayscale) images and a custom fully connected layer to output CAPTCHA characters.
- Feature Extractor:
models.resnet18(). - Input Adaptation:
torch.nn.Conv2d(1, 64, kernel_size=(50, 160)...). - Output Layer:
torch.nn.Linear(in_features, out_features=chars_len * max_capcha). - Post-Processing in Graph:
x.view(-1, max_capcha, chars_len)to organize output into character-wise probabilities.
Sources: train_model/core.py:68-81
Deployment and Inference
Once exported, the model.onnx file is integrated into the mb_capcha_ocr package.
Package Integration
The setup.py script includes the ONNX file as package data, ensuring it is distributed when the library is installed via pip.
package_data={
'mb_capcha_ocr': ['model.onnx'],
},
Sources: setup.py:24-26
Inference Execution
The predict.py module uses onnxruntime to load the exported file. It maps the input image to the "input" name defined during the export and processes the "output" tensor to extract text.
The sequence illustrates how the exported ONNX model is utilized by the inference engine.
Sources: mb_capcha_ocr/predict.py:15-32, README.md:31-38
Dependencies for Exporting
A specific set of requirements is needed for the export phase, which differs from the lightweight requirements for using the exported model.
| Phase | Key Dependencies |
|---|---|
| Exporting (Training) | torch, torchvision, onnx, onnxscript |
| Using (Inference) | onnxruntime, numpy, Pillow |
Sources: train_model/train_requirements.txt:1-7, requirements.txt:1-5
Exporting to ONNX is a critical step in this project that transforms a PyTorch-based training workflow into a portable, high-performance inference module. This allows the OCR capabilities to be easily consumed by other applications with minimal dependency overhead.