PyTorch Neural Network Architecture

Relevant source files

The following files were used as context for generating this wiki page:

• train_model/core.py
• train_model/train.py
• mb_capcha_ocr/predict.py
• README.md
• setup.py

The PyTorch Neural Network Architecture in this project is designed for Optical Character Recognition (OCR), specifically targeting MBBank CAPTCHA images. The system leverages a modified ResNet-18 backbone to extract features from grayscale images and map them to a multi-label classification output representing a sequence of characters.

The architecture facilitates a complete pipeline from raw image data processing and model training to deployment via ONNX exportation. It utilizes a custom dataset loader to handle CAPTCHA-specific labeling and an inference class for production-ready predictions.
Sources: README.md:1-5, train_model/core.py:100-112

Model Structure

The core of the system is the OcrModel class, which inherits from torch.nn.Module. It utilizes a transfer learning approach by modifying a standard ResNet-18 architecture to suit the specific dimensions and requirements of grayscale CAPTCHA recognition.
Sources: train_model/core.py:100-101

Architecture Components

The model modifies the standard ResNet-18 in two primary ways:

1. Input Layer: The initial convolutional layer (conv1) is replaced to accept a single-channel (grayscale) input instead of the default three-channel RGB. It uses a specific kernel size of (50, 160) and a stride of (2, 2).
2. Output Layer: The fully connected layer (fc) is modified to output a flattened vector of size chars_len * max_capcha. This allows the model to predict multiple characters simultaneously.

Component	Specification	Description
Backbone	ResNet-18	Feature extraction base
Input Channel	1 (Grayscale)	Optimized for single-channel CAPTCHA images
Output Features	chars_len * max_capcha	62 possible characters * 6 positions = 372 features
Reshape Target	(-1, 6, 62)	Reshaped for multi-label interpretation

Sources: train_model/core.py:103-112, train_model/train.py:24-27

Data Flow Diagram

The following diagram illustrates the transformation of an image through the neural network layers during a forward pass.

The forward pass reshapes the linear output into a 3D tensor where the second dimension represents the character position in the CAPTCHA and the third dimension represents the character class probability.
Sources: train_model/core.py:108-112

Data Processing and Loading

The DatasetLoader class handles the ingestion of CAPTCHA images and their corresponding labels. Labels are extracted directly from the filenames and converted into one-hot encoded vectors.

Label Encoding

Labels are derived from a character set containing digits and both uppercase and lowercase ASCII letters.

• Character Set: 0-9, a-z, A-Z (Sorted alphabetically).
• Max Length: Fixed at 6 characters.
• Encoding: Each character is converted to a one-hot vector of length 62.

Training Pipeline

The training process uses MultiLabelSoftMarginLoss to optimize the model against the multi-character output.

Parameter	Value	Source
Optimizer	Adam	train_model/train.py:44
Learning Rate	0.001	train_model/train.py:44
Loss Function	MultiLabelSoftMarginLoss	train_model/train.py:43
Batch Size	64	train_model/train.py:35
Epochs	100	train_model/train.py:47

Prediction and Inference

After training, the model is exported to ONNX format for efficient inference. The mb_capcha_ocr.predict.OcrModel class handles the processing of images using onnxruntime.

Inference Sequence

The prediction logic mirrors the training transformations to ensure consistency.

Conclusion

The PyTorch Neural Network Architecture implemented in this repository provides a streamlined approach to CAPTCHA recognition by adapting a proven image classification backbone (ResNet-18) for sequence-based character prediction. By standardizing the input to 160x50 grayscale images and outputting a reshaped multi-label tensor, the system achieves a balance between architectural complexity and specialized OCR functionality. The integration of ONNX ensures that the trained PyTorch models can be deployed in production environments with minimal overhead.