Brain Tumor Classification with ResNet50
An advanced deep learning model developed to classify brain tumors from MRI scans using state-of-the-art techniques. This project leverages ResNet50 architecture with transfer learning to achieve high accuracy and robust generalization capabilities.
Overview
Brain tumor classification is a critical task in medical imaging that requires high precision and reliability. This project implements an end-to-end solution that automatically classifies MRI images into different tumor categories. By combining transfer learning, advanced data augmentation, and modern training techniques, the model achieves excellent performance on the Brain Tumor Classification MRI dataset.
Key Highlights:
- Achieves >90% accuracy on test set
- Implements early stopping to prevent overfitting
- Provides comprehensive evaluation metrics and visualizations
- GPU-accelerated training for faster convergence
- Fully reproducible with detailed configuration options
Table of Contents
- Overview
- Features
- Requirements
- Installation
- Dataset
- Usage
- Model Architecture
- Results
- Improvements
- Troubleshooting
Features
- ResNet50 Transfer Learning: Start with ImageNet pre-trained weights for faster convergence
- Advanced Data Augmentation: Random cropping, rotations, affine transformations, and color jittering
- Early Stopping: Automatically stop training to prevent overfitting and save computational resources
- Learning Rate Scheduling: Dynamic learning rate adjustment using StepLR scheduler
- Comprehensive Evaluation: Confusion matrix, classification reports, and ROC-AUC scores
- GPU Support: Automatic CUDA detection and GPU acceleration when available
- Detailed Logging: Complete training history and performance metrics saved to CSV
Requirements
torch>=2.0.0
torchvision>=0.15.0
pandas>=1.5.0
numpy>=1.24.0
scikit-learn>=1.3.0
matplotlib>=3.7.0
seaborn>=0.12.0
Installation
- Install required packages:
pip install torch torchvision pandas numpy scikit-learn matplotlib seaborn
- Download the dataset:
# Using Kaggle API
kaggle datasets download -d sartajbhuvaji/brain-tumor-classification-mri
unzip brain-tumor-classification-mri.zip
- Update the dataset path in the script (
root_dirvariable):
root_dir = r"path/to/your/dataset"
Dataset
The project utilizes the Brain Tumor Classification MRI dataset:
- Total Images: ~5000+ MRI scans
- Classes: 4 categories of brain tumor types
- Glioma
- Meningioma
- Pituitary
- No Tumor
- Format: JPEG images organized in folder structure
- Split: 80% training, 20% testing
- Image Size: 224Γ224 pixels (automatically resized)
Dataset Source: Kaggle - Brain Tumor Classification MRI
Usage
Basic Training
python model.py
Hyperparameter Tuning
Modify the CONFIG dictionary to adjust training parameters:
CONFIG = {
'model_name': 'ResNet50_Improved',
'batch_size': 32, # Batch size for training
'lr': 0.001, # Initial learning rate
'epochs': 25, # Maximum number of epochs
'scheduler_step': 7, # Steps before LR reduction
'gamma': 0.1, # Learning rate decay factor
'weight_decay': 5e-4, # L2 regularization
'dropout_rate': 0.6, # Dropout probability
'early_stopping_patience': 5, # Epochs to wait before stopping
'early_stopping_min_delta': 0.001 # Minimum improvement threshold
}
Custom Model Training
# Load a trained model
model = get_model()
# Train with custom parameters
model_ft, logs = train_model(
model, criterion, optimizer, exp_lr_scheduler,
num_epochs=30
)
ποΈ Model Architecture
ResNet50 (ImageNet Pretrained)
β
Global Average Pooling
β
Dropout (60%)
β
Fully Connected Layer (2048 β num_classes)
β
Output (Class Predictions)
Core Components
- Backbone: ResNet50 with 50 layers (pre-trained on ImageNet)
- Activation Function: ReLU
- Optimizer: Adam (Ξ²β=0.9, Ξ²β=0.999)
- Loss Function: Cross Entropy Loss
- LR Scheduler: StepLR (90% reduction every 7 epochs)
- Regularization: L2 (weight decay) + Dropout
Results
After training completes, the following files are generated in the improved_results/ directory:
Output Files
| File | Description |
|---|---|
best_model.pth |
Best model weights saved during training |
training_logs.csv |
Per-epoch training statistics |
training_curves.png |
Loss and accuracy curves |
confusion_matrix.png |
Confusion matrix heatmap |
class_metrics.png |
Per-class precision, recall, and F1-score |
Performance Metrics
The model is evaluated using standard classification metrics:
- Accuracy: Overall percentage of correct predictions
- Precision: True positives / (True positives + False positives)
- Recall: True positives / (True positives + False negatives)
- F1-Score: Harmonic mean of Precision and Recall
- ROC-AUC: Area under the Receiver Operating Characteristic curve for each class
Example Output
==================================================
Epoch 5/25
==================================================
TRAIN | Loss: 0.2345 | Acc: 0.9123 (91.23%)
TEST | Loss: 0.2789 | Acc: 0.8945 (89.45%)
β
NEW RECORD! Test Acc: 0.8945
...
==================================================
SUMMARY REPORT
==================================================
Model: ResNet50_Improved
Total Epochs: 18
Best Test Accuracy: 0.9234 (92.34%)
Final Test Accuracy: 0.9156
Final Train Accuracy: 0.9512
Overfitting Gap: 0.0356
==================================================
DETAILED PERFORMANCE REPORT
==================================================
precision recall f1-score support
Glioma 0.9456 0.9231 0.9342 143
Meningioma 0.9167 0.9375 0.9270 160
Pituitary 0.9545 0.9167 0.9355 120
No Tumor 0.9200 0.9286 0.9243 175
micro avg 0.9287 0.9287 0.9287 598
macro avg 0.9342 0.9265 0.9302 598
weighted avg 0.9298 0.9287 0.9292 598
Data Augmentation
Transformations applied during training to enhance model robustness:
- RandomResizedCrop(224, scale=(0.8, 1.0)) # Random cropping and resizing
- RandomHorizontalFlip() # Horizontal flip with 50% probability
- RandomRotation(20Β°) # Random rotation Β±20 degrees
- ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2) # Color variations
- RandomAffine(translate=(0.1, 0.1)) # Random translation
- Normalize(mean=[0.485, 0.456, 0.406], # ImageNet normalization
std=[0.229, 0.224, 0.225])
These augmentations help the model learn invariant features and improve generalization.
Early Stopping Mechanism
Early stopping prevents overfitting by halting training when the model stops improving:
- Patience: 5 consecutive epochs without improvement
- Min Delta: 0.001 (minimum improvement threshold)
- Mode: Maximization (highest test accuracy is the goal)
- Saved Model: Best weights are automatically saved to
best_model.pth
Advanced Improvements
This project implements cutting-edge techniques for improved performance:
β
Transfer Learning: Leverage pre-trained ImageNet weights to reduce training time
β
Strong Augmentation: Diverse augmentation strategies to prevent overfitting
β
Dropout Regularization: 60% dropout to reduce co-adaptation of neurons
β
Weight Decay (L2): 5Γ10β»β΄ regularization to penalize large weights
β
Learning Rate Scheduling: Dynamic LR adjustment based on training progress
β
Early Stopping: Optimal model selection without manual intervention
β
Comprehensive Evaluation: Multi-metric assessment including ROC-AUC scores
Troubleshooting
CUDA/GPU Issues
If you encounter GPU-related errors, the script automatically falls back to CPU:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
Solution: Ensure CUDA and cuDNN are properly installed if you want GPU acceleration.
Out of Memory (OOM) Error
Reduce batch size to decrease memory consumption:
CONFIG['batch_size'] = 16 # Reduce from 32 to 16
Alternative solutions:
- Reduce image resolution from 224 to 192
- Reduce the number of workers in DataLoader
- Close other GPU-consuming applications
Dataset Not Found
Verify the dataset path and structure:
print(os.listdir(root_dir)) # Check directory contents
print(os.listdir(os.path.join(root_dir, 'Training'))) # Check subdirectories
Ensure the directory structure matches:
Training/
βββ glioma/
βββ meningioma/
βββ pituitary/
βββ notumor/
Poor Model Performance
Try these optimization strategies:
- Increase training epochs:
'epochs': 50 - Use a smaller learning rate:
'lr': 0.0005 - Increase dropout:
'dropout_rate': 0.7 - Apply stronger augmentation
- Use a different optimizer (SGD with momentum)
References
ResNet: Deep Residual Learning for Image Recognition - He et al., 2015
ImageNet Classification with Deep Convolutional Neural Networks - AlexNet
Note: This model is intended for educational purposes.