Performance Optimization
Relevant source files
The following files were used as context for generating this wiki page:
Performance optimization in the MBBank CAPTCHA OCR project is achieved through a combination of efficient model architecture, optimized runtime inference, and controlled data management during training. The system transitions from a heavy PyTorch training environment to a lightweight, high-performance ONNX runtime for production use.
The optimization strategy focuses on reducing the computational footprint of inference while maintaining high accuracy for CAPTCHA recognition. This is facilitated by using standardized image preprocessing and the deployment of a serialized model format.
Inference Optimization via ONNX
The project optimizes the inference phase by exporting the trained PyTorch model to the Open Neural Network Exchange (ONNX) format. This allows the application to utilize onnxruntime for predictions, which is generally faster and requires fewer dependencies than the full PyTorch library in a production environment.
Model Deployment Flow
The following diagram illustrates the transition from training to optimized inference:
The model is exported using torch.onnx.export with defined input and output names to ensure compatibility with the inference session.
Sources: train_model/train.py:56-57, mb_capcha_ocr/predict.py:16-17
Inference Execution
The OcrModel class in the prediction module initializes an onnxruntime.InferenceSession. This session is configured to load the serialized model.onnx file, providing a dedicated environment for executing the model graph efficiently.
Sources: mb_capcha_ocr/predict.py:15-17, setup.py:22-24
| Optimization Component | Implementation | Benefit |
|---|---|---|
| Runtime | onnxruntime | Reduced overhead compared to torch during inference. |
| Model Format | .onnx | Portability and cross-platform performance consistency. |
| Initialization | InferenceSession | Pre-loads the model graph into memory for rapid reuse. |
Sources: mb_capcha_ocr/predict.py:1-17, requirements.txt:1
Architectural Efficiency
The underlying architecture utilizes a modified ResNet-18 backbone. Performance is optimized by tailoring the network layers to the specific dimensions and characteristics of the CAPTCHA images.
Layer Modifications
To improve efficiency for small grayscale images, the first convolutional layer and the fully connected layer are customized:
- Conv2d Optimization: The initial layer is adjusted to handle single-channel (grayscale) input with a specific kernel size of (50, 160) and stride of (2, 2).
- Linear Layer Output: The final layer is mapped directly to
chars_len * max_capcha(372 outputs), eliminating redundant calculations for unrelated character sets.
Sources: train_model/core.py:76-78, train_model/core.py:14-16
The output is reshaped to group predictions by character position, allowing for vectorized decoding using argmax.
Sources: train_model/core.py:81-83, mb_capcha_ocr/predict.py:24-26
Data Handling and Training Performance
During the training phase, several techniques are employed to manage performance and prevent resource exhaustion.
Dataset Limiting
To prevent overfitting and maintain manageable training times, the DatasetLoader implements a hard limit on the number of images processed from the source directory.
- Cap: 62,500 images.
- Split: 10% of the capped data is reserved for testing.
Sources: train_model/core.py:42-47
Parallel Processing
The training script utilizes torch.utils.data.DataLoader with batching to optimize GPU/CPU utilization.
- Batch Size: 64.
- Shuffle: Enabled for training to improve convergence speed.
Sources: train_model/train.py:29-30
Image Preprocessing Pipeline
Uniform image processing ensures that both training and inference spend minimal time on transformation. Images are converted to "L" (Grayscale) and resized to 160x50 pixels.
# Inference Preprocessing Optimization
img = img_data.convert("L")
img = img.resize((160, 50))
img = np.array(img, dtype=np.float32) / 255.0
Sources: mb_capcha_ocr/predict.py:20-22, train_model/core.py:64-68
Conclusion
Performance optimization in the mbbank-capcha-ocr project is centered on the strategic transition from PyTorch training to ONNX-based inference. By combining ResNet-18 architectural refinements with efficient data batching and the lightweight onnxruntime, the system achieves a high-throughput capability suitable for real-time CAPTCHA recognition.