Troubleshooting Training Issues
Relevant source files
The following files were used as context for generating this wiki page:
This page provides a technical overview of potential issues and debugging strategies for the MBBank CAPTCHA OCR training pipeline. The training process involves a PyTorch-based ResNet18 architecture modified for multi-label classification, utilizing a specific dataset format and image preprocessing workflow.
Effective troubleshooting requires understanding the interaction between the DatasetLoader, the model architecture defined in core.py, and the execution loop in train.py. Issues typically arise from data formatting, environmental dependencies, or convergence failures during the training epochs.
Data Pipeline and Preprocessing
The training system relies on a specific image naming convention and directory structure. The DatasetLoader class parses filenames to extract labels, which are then converted into one-hot encoded vectors.
Dataset Structure Requirements
Training and testing images must be located in a dataset/ directory. The filename itself serves as the ground truth label.
- Format:
{label}.(png|jpg|jpeg) - Constraint: The
DatasetLoaderimposes a hard limit of 62,500 images to prevent overfitting. - Split: The dataset is automatically split, with 10% (or at least 1 image) allocated for testing.
Sources: README.md:31-33, train_model/train.py:23, train_model/core.py:38-48
The diagram above illustrates the data flow within the DatasetLoader class.
Sources: train_model/core.py:32-60
Image Transformation Issues
The model expects a fixed input size of 160x50 pixels. Deviations in the transforms.Resize parameters or failure to convert images to Grayscale ("L") can lead to tensor dimension mismatches or performance degradation.
| Parameter | Value | Description | Source |
|---|---|---|---|
| Target Height | 50 | Vertical resolution for ResNet input | train.py:26 |
| Target Width | 160 | Horizontal resolution for ResNet input | train.py:26 |
| Color Mode | "L" | Images are converted to Grayscale | core.py:65 |
| Normalization | ToTensor | Scales pixels to [0.0, 1.0] | train.py:27 |
Model Architecture and Loss Convergence
The OCR model uses a modified ResNet18. If training loss does not decrease, it may be due to the architecture's specific configuration for multi-label classification.
ResNet18 Modifications
The standard ResNet18 is modified in two critical ways:
- Input Layer: The first convolutional layer (
conv1) is redefined to accept a single-channel (grayscale) input with a specific kernel size of (50, 160). - Output Layer: The fully connected layer (
fc) is modified to output a flat vector of sizechars_len * max_capcha(typically $62 \times 6$).
Sources: train_model/core.py:73-81
Loss Function and Optimization
The trainer uses MultiLabelSoftMarginLoss. This is specifically chosen because the model treats the CAPTCHA as a sequence of independent multi-class classification tasks flattened into one output.
model = OcrModel().to(device)
loss_func = torch.nn.MultiLabelSoftMarginLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
Sources: train_model/train.py:38-40
Training Flow and Monitoring
Training progress is monitored via standard output and a generated loss.png plot.
- Epochs: Default is set to 100.
- Batch Size: 64 images per step.
- Visualization:
matplotlibis used to plotloss_valuesagainst iterations.
This sequence shows the standard optimization loop.
Sources: train_model/train.py:43-55
Environment and Dependency Issues
Failures during training often stem from missing dependencies or incorrect CUDA configurations.
Requirements for Training
The training environment requires a larger set of dependencies than the inference environment, including torch, torchvision, and onnx.
| Package | Purpose | File |
|---|---|---|
torch | Core ML framework | train_requirements.txt:1 |
torchvision | Image transforms and ResNet base | train_requirements.txt:2 |
onnx | Model export | train_requirements.txt:6 |
matplotlib | Troubleshooting/Loss plotting | train_requirements.txt:5 |
pillow | Image loading | train_requirements.txt:3 |
CUDA vs. CPU Execution
The script automatically detects the availability of a GPU. If the training is excessively slow, users should verify if torch.cuda.is_available() returns true.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
Sources: train_model/train.py:20
Export and Prediction Validation
A successful training run concludes with the export of the model to ONNX format. Troubleshooting must ensure the dimensions used during training match the expectations of the inference session in predict.py.
Dimension Consistency
The OcrModel.forward method reshapes output to (-1, max_capcha, chars_len). The predict.py module expects this specific 3D shape during np.argmax.
| Feature | Training Implementation | Prediction Implementation |
|---|---|---|
| Resizing | (50, 160) | (160, 50) (PIL format) |
| Input Shape | [Batch, 1, 50, 160] | [1, 1, 50, 160] |
| Argmax Axis | dim=2 | axis=2 |
Sources: train_model/train.py:73, mb_capcha_ocr/predict.py:22-31, train_model/core.py:85
Common Export Errors
The model is exported using torch.onnx.export. If modal.onnx (the filename specified in train.py) is missing, ensure the script has write permissions in the directory. Note that the setup.py and predict.py expect the file to be named model.onnx.
Sources: train_model/train.py:60, setup.py:26, mb_capcha_ocr/predict.py:16
Summary
Troubleshooting training in the mbbank-capcha-ocr project focuses on ensuring data integrity in the dataset/ folder, verifying the modified ResNet18 architecture dimensions, and maintaining environment parity with train_requirements.txt. Monitoring the loss.png output is the primary method for identifying convergence issues during the 100-epoch training cycle.