Training Environment Setup

Relevant source files

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

• README.md
• train_model/train_requirements.txt
• train_model/train.py
• train_model/core.py
• requirements.txt
• setup.py

The training environment for the MBBank CAPCHA OCR project is a specialized configuration designed to facilitate the development, training, and exportation of a ResNet18-based deep learning model. This environment focuses on processing image datasets, calculating multi-label loss, and generating an ONNX format model for production use.

Sources: README.md:1-12, train_model/train.py:10-25, train_model/core.py:84-90

System Requirements and Dependencies

The environment requires Python 3.x and is split into runtime dependencies and extended training dependencies. While the core library only requires lightweight runtimes like onnxruntime, the training environment necessitates heavy-duty frameworks for tensor computation and automatic differentiation.

Core Training Dependencies

Dependency	Purpose	Source
torch	Main deep learning framework for model architecture and training	train_model/train_requirements.txt:1
torchvision	Provides the ResNet18 base model and image transformation utilities	train_model/train_requirements.txt:2
pillow	Image loading and manipulation (grayscale conversion)	train_model/train_requirements.txt:3
numpy	Numerical operations and array handling	train_model/train_requirements.txt:4
matplotlib	Visualization of training loss curves	train_model/train_requirements.txt:5
onnx / onnxscript	Utilities for exporting the trained model to ONNX format	train_model/train_requirements.txt:6-7

Sources: train_model/train_requirements.txt:1-7, README.md:46-52

Environment Initialization

Setting up the training environment involves creating an isolated virtual environment and installing specific requirements located within the train_model directory.

The sequence above ensures that the training scripts have access to the necessary libraries without conflicting with the global Python installation.
Sources: README.md:14-30

Data Pipeline Architecture

The training environment expects a specific directory structure and data format. Images must be placed in a dataset/ folder with filenames serving as their labels.

Dataset Configuration

The DatasetLoader class manages the data flow, implementing a 90/10 split for training and testing by default. It also enforces a size limit of 62,500 images to prevent overfitting.

Training Logic and Hardware Acceleration

The environment is configured to automatically detect and utilize NVIDIA GPUs via CUDA if available, falling back to CPU otherwise.

Key Components

• Model Architecture: A modified ResNet18. The first convolutional layer is adapted to accept 1-channel grayscale input, and the fully connected layer is reshaped to output chars_len * max_capcha (62 characters * 6 positions).
• Loss Function: MultiLabelSoftMarginLoss is used to handle the multi-character prediction task.
• Optimizer: Adam optimizer with a learning rate of 0.001.

Model Export and Validation

Upon completion of the training cycles, the environment performs two final tasks: generating a visual loss report and exporting the model for production.

1. Loss Visualization: Uses matplotlib to plot loss_values and saves the result as loss.png.
2. ONNX Export: Converts the PyTorch model to an ONNX graph using a dummy input tensor of shape (1, 1, 50, 160). The resulting file is named modal.onnx.
3. Accuracy Testing: Iterates through the test_dl to compare the argmax of predicted characters against the actual image labels.

Sources: train_model/train.py:62-85

Summary

The training environment setup provides a robust pipeline for CAPCHA recognition, transitioning from raw images in a dataset/ directory to a deployable ONNX model. By leveraging ResNet18 and custom data loaders, it automates the transformation of grayscale imagery into character sequences.