The Training Loop
Relevant source files
The following files were used as context for generating this wiki page:
The training loop is a centralized process within the mbbank-capcha-ocr project designed to optimize a Deep Learning model for CAPTCHA recognition. It orchestrates data ingestion, model optimization using PyTorch, and final export to the ONNX format for efficient inference. The process specifically targets the recognition of alphanumeric characters typically found in MBBank CAPTCHAs, employing a modified ResNet18 architecture.
Sources: README.md:1-10, train_model/train.py:1-25, train_model/core.py:46-55
Data Preparation and Loading
The training loop begins by initializing a specialized data pipeline. The system expects images in a dataset/ directory where filenames act as the ground-truth labels (e.g., label.png). The DatasetLoader class handles the discovery, shuffling, and splitting of these files into training and validation sets.
Dataset Configuration
The dataset is processed with specific constraints to ensure model stability and performance:
- Resizing: Images are normalized to a resolution of 160x50 pixels.
- Grayscale Conversion: Images are converted to "L" (8-bit pixels, black and white) mode.
- One-Hot Encoding: Labels are converted into one-hot encoded vectors based on a character set of digits and letters.
- Dataset Split: 10% of the data is reserved for testing, while the remaining 90% is used for training.
Sources: train_model/train.py:27-38, train_model/core.py:27-44
The diagram above illustrates the data flow from raw image files to the PyTorch DataLoader. Sources: train_model/train.py:32-38, train_model/core.py:46-77
Model Architecture
The core of the training loop is the OcrModel, which leverages a modified ResNet18 backbone. The architecture is adapted specifically for 1-channel grayscale input and fixed-length CAPTCHA output.
| Component | Modification | Description |
|---|---|---|
| Input Layer | conv1 | Modified to accept 1 input channel (grayscale) with kernel size (50, 160). |
| Backbone | resnet18 | Standard ResNet18 feature extractor. |
| Output Layer | fc (Linear) | Linear layer mapping to chars_len * max_capcha (62 characters * 6 positions). |
| Reshape | forward() | Output is reshaped to (-1, 6, 62) for multi-label classification. |
Sources: train_model/core.py:80-92
Optimization Loop
The training process executes for 100 epochs, utilizing the Adam optimizer and a multi-label margin loss function.
Training Logic
- Zero Gradients: The optimizer clears previous gradients.
- Forward Pass: The model predicts the character sequence for a batch of images.
- Loss Calculation:
MultiLabelSoftMarginLossmeasures the difference between predictions and one-hot encoded labels. - Backward Pass: Gradients are computed via backpropagation.
- Step: Weights are updated based on the calculated gradients.
Sources: train_model/train.py:40-60
The sequence diagram represents the iterative weight update process during the training phase. Sources: train_model/train.py:47-60
Evaluation and Export
Following training, the loop performs two final actions: performance evaluation and model serialization.
Model Evaluation
The model is switched to evaluation mode (model.eval()). It processes the test dataset, comparing the argmax of predicted tensors against the actual text labels to calculate accuracy.
ONNX Export
The trained PyTorch model is exported to the Open Neural Network Exchange (ONNX) format. This allows the model to be used for inference in environments without a full PyTorch installation, using onnxruntime.
# Exporting the model
torch_input = torch.randn(1, 1, 50, 160).to(device)
torch.onnx.export(model, (torch_input,), f"modal.onnx",
input_names=['input'],
output_names=['output'])
Sources: train_model/train.py:62-63, train_model/train.py:73-88
Summary
The Training Loop in mbbank-capcha-ocr provides a robust pipeline for converting labeled CAPTCHA images into a portable ONNX model. By utilizing a ResNet18-based architecture and a MultiLabelSoftMarginLoss approach, it effectively learns to map visual patterns to a 6-character alphanumeric sequence. The resulting model.onnx is then used by the predict.py module for production OCR tasks.