| # 🎯 Classification Training Guide |
|
|
| Complete guide for fine-tuning the BYOL pre-trained model for multi-label classification. |
|
|
| ## 📋 Overview |
|
|
| After BYOL pre-training completes, you can fine-tune the model for classification using the `train_classification.py` script. This approach: |
|
|
| 1. **Loads the BYOL checkpoint** with learned representations |
| 2. **Freezes the backbone** initially (optional) to prevent overwriting good features |
| 3. **Fine-tunes the classification head** with a higher learning rate |
| 4. **Gradually unfreezes** the backbone for end-to-end fine-tuning |
|
|
| ## 🗂️ Data Preparation |
|
|
| ### CSV Format |
| Create train/validation CSV files with this format: |
|
|
| ```csv |
| tile_path,mass,calcification,architectural_distortion,asymmetry,normal,benign,malignant,birads_2,birads_3,birads_4 |
| patient1_tile_001.png,1,0,0,0,0,1,0,0,1,0 |
| patient1_tile_002.png,0,1,0,0,0,0,1,0,0,1 |
| patient2_tile_001.png,0,0,0,0,1,1,0,1,0,0 |
| ... |
| ``` |
|
|
| **Requirements:** |
| - `tile_path`: Relative path to tile image |
| - **Class columns**: Binary labels (0/1) for each class |
| - **Multi-label support**: Each image can have multiple classes = 1 |
|
|
| ### Directory Structure |
| ``` |
| your_project/ |
| ├── tiles/ # Directory containing tile images |
| │ ├── patient1_tile_001.png |
| │ ├── patient1_tile_002.png |
| │ └── ... |
| ├── train_labels.csv # Training labels |
| ├── val_labels.csv # Validation labels |
| └── mammogram_byol_best.pth # BYOL checkpoint |
| ``` |
|
|
| ## 🚀 Quick Start |
|
|
| ### 1. Basic Classification Training |
|
|
| ```bash |
| python train_classification.py \ |
| --byol_checkpoint ./mammogram_byol_best.pth \ |
| --train_csv ./train_labels.csv \ |
| --val_csv ./val_labels.csv \ |
| --tiles_dir ./tiles \ |
| --class_names mass calcification architectural_distortion asymmetry normal benign malignant birads_2 birads_3 birads_4 \ |
| --output_dir ./classification_results |
| ``` |
|
|
| ### 2. With Custom Configuration |
|
|
| ```bash |
| python train_classification.py \ |
| --byol_checkpoint ./mammogram_byol_best.pth \ |
| --train_csv ./train_labels.csv \ |
| --val_csv ./val_labels.csv \ |
| --tiles_dir ./tiles \ |
| --class_names mass calcification normal \ |
| --config ./classification_config.json \ |
| --output_dir ./classification_results \ |
| --wandb_project my-mammogram-classification |
| ``` |
|
|
| ### 3. Quick Testing (Limited Dataset) |
|
|
| ```bash |
| python train_classification.py \ |
| --byol_checkpoint ./mammogram_byol_best.pth \ |
| --train_csv ./train_labels.csv \ |
| --val_csv ./val_labels.csv \ |
| --tiles_dir ./tiles \ |
| --class_names mass calcification normal \ |
| --max_samples 1000 \ |
| --output_dir ./test_results |
| ``` |
|
|
| ## ⚙️ Configuration Options |
|
|
| ### Key Parameters |
|
|
| | Parameter | Default | Description | |
| |-----------|---------|-------------| |
| | `batch_size` | 32 | Batch size for training | |
| | `epochs` | 50 | Number of training epochs | |
| | `lr_backbone` | 1e-5 | Learning rate for pre-trained backbone | |
| | `lr_head` | 1e-3 | Learning rate for classification head | |
| | `freeze_backbone_epochs` | 10 | Epochs to freeze backbone (0 = never freeze) | |
| | `label_smoothing` | 0.1 | Label smoothing for regularization | |
| | `gradient_clip` | 1.0 | Gradient clipping max norm | |
|
|
| ### Custom Configuration File |
|
|
| Create `my_config.json`: |
| ```json |
| { |
| "batch_size": 64, |
| "epochs": 100, |
| "lr_backbone": 5e-6, |
| "lr_head": 2e-3, |
| "freeze_backbone_epochs": 20, |
| "label_smoothing": 0.2, |
| "weight_decay": 1e-3 |
| } |
| ``` |
|
|
| ## 📊 Expected Training Process |
|
|
| ### Phase 1: Backbone Frozen (Epochs 1-10) |
| ``` |
| 🧊 Epoch 1: Backbone frozen (training only classification head) |
| Epoch 1/50: |
| Train Loss: 0.6234 |
| Val Loss: 0.5891 |
| Mean AUC: 0.7123 |
| Mean AP: 0.6894 |
| Exact Match: 0.4512 |
| ✅ New best model saved (AUC: 0.7123) |
| ``` |
|
|
| ### Phase 2: End-to-End Fine-tuning (Epochs 11-50) |
| ``` |
| Epoch 15/50: |
| Train Loss: 0.3456 |
| Val Loss: 0.3891 |
| Mean AUC: 0.8567 |
| Mean AP: 0.8234 |
| Exact Match: 0.6789 |
| ✅ New best model saved (AUC: 0.8567) |
| ``` |
|
|
| ## 🔍 Making Predictions |
|
|
| ### Single Image Inference |
|
|
| ```bash |
| python inference_classification.py \ |
| --model_path ./classification_results/best_classification_model.pth \ |
| --image_path ./test_image.png \ |
| --threshold 0.5 |
| ``` |
|
|
| **Output:** |
| ``` |
| 📸 Image 1: test_image.png |
| 🏆 Top prediction: mass (0.847) |
| 📊 All probabilities: |
| ✅ mass : 0.847 |
| ❌ calcification : 0.234 |
| ❌ normal : 0.123 |
| ❌ architectural_distortion: 0.089 |
| ``` |
|
|
| ### Batch Inference |
|
|
| ```bash |
| python inference_classification.py \ |
| --model_path ./classification_results/best_classification_model.pth \ |
| --images_dir ./test_images \ |
| --output_json ./predictions.json \ |
| --batch_size 64 |
| ``` |
|
|
| ### Programmatic Usage |
|
|
| ```python |
| import torch |
| from train_byol_mammo import MammogramBYOL |
| from inference_classification import load_classification_model, create_inference_transform |
| |
| # Load model |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| model, class_names, config = load_classification_model( |
| "./classification_results/best_classification_model.pth", device |
| ) |
| |
| # Make prediction |
| transform = create_inference_transform() |
| image = Image.open("test.png").convert('RGB') |
| input_tensor = transform(image).unsqueeze(0).to(device) |
| |
| with torch.no_grad(): |
| logits = model.classify(input_tensor) |
| probabilities = torch.sigmoid(logits).cpu().numpy()[0] |
| |
| # Get results |
| for i, class_name in enumerate(class_names): |
| print(f"{class_name}: {probabilities[i]:.3f}") |
| ``` |
|
|
| ## 📈 Monitoring Training |
|
|
| ### Weights & Biases Integration |
|
|
| The script automatically logs to W&B: |
| - Training/validation loss curves |
| - Per-class AUC and Average Precision |
| - Learning rate schedules |
| - Model hyperparameters |
|
|
| ### Metrics Explained |
|
|
| - **AUC (Area Under Curve)**: Measures ranking quality (0-1, higher better) |
| - **AP (Average Precision)**: Summarizes precision-recall curve (0-1, higher better) |
| - **Exact Match Accuracy**: Percentage where ALL labels are predicted correctly |
| - **Per-Class Accuracy**: Binary accuracy for each individual class |
|
|
| ## 💾 Output Files |
|
|
| Training creates: |
| ``` |
| classification_results/ |
| ├── best_classification_model.pth # Best model by validation AUC |
| ├── final_classification_model.pth # Final model after all epochs |
| ├── classification_epoch_10.pth # Periodic checkpoints |
| ├── classification_epoch_20.pth |
| └── ... |
| ``` |
|
|
| Each checkpoint contains: |
| - Model state dict |
| - Optimizer state |
| - Training configuration |
| - Class names |
| - Validation metrics |
|
|
| ## 🛠️ Advanced Usage |
|
|
| ### Custom Loss Functions |
|
|
| For imbalanced datasets, modify the loss function: |
|
|
| ```python |
| # Calculate positive weights for each class |
| pos_counts = df[class_names].sum() |
| neg_counts = len(df) - pos_counts |
| pos_weight = torch.tensor(neg_counts / pos_counts).to(device) |
| |
| criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight) |
| ``` |
|
|
| ### Transfer Learning Strategies |
|
|
| 1. **Conservative**: Freeze backbone for many epochs, low backbone LR |
| - `freeze_backbone_epochs = 20` |
| - `lr_backbone = 1e-6` |
|
|
| 2. **Aggressive**: Unfreeze early, higher backbone LR |
| - `freeze_backbone_epochs = 5` |
| - `lr_backbone = 1e-4` |
|
|
| 3. **Progressive**: Gradually unfreeze layers (requires code modification) |
|
|
| ### Multi-GPU Training |
|
|
| For multiple GPUs, wrap the model: |
| ```python |
| if torch.cuda.device_count() > 1: |
| model = nn.DataParallel(model) |
| ``` |
|
|
| ## ⚠️ Troubleshooting |
|
|
| ### Common Issues |
|
|
| **Low Validation Performance:** |
| - Increase `freeze_backbone_epochs` to 15-20 |
| - Reduce `lr_backbone` to 5e-6 or 1e-6 |
| - Check for data leakage between train/val sets |
|
|
| **Overfitting:** |
| - Increase `label_smoothing` to 0.2-0.3 |
| - Add more dropout (modify model architecture) |
| - Reduce learning rates |
| - Use early stopping |
|
|
| **Memory Issues:** |
| - Reduce `batch_size` to 16 or 8 |
| - Reduce `num_workers` to 4 |
| - Use gradient checkpointing (requires code modification) |
|
|
| **Class Imbalance:** |
| - Use `pos_weight` in loss function |
| - Focus on per-class AUC rather than accuracy |
| - Consider focal loss for extreme imbalance |
|
|
| ## 🎯 Best Practices |
|
|
| 1. **Start Conservative**: Use default settings first |
| 2. **Monitor Per-Class Metrics**: Some classes may need special attention |
| 3. **Validate Data**: Ensure no train/val overlap |
| 4. **Checkpoint Often**: Training can be interrupted |
| 5. **Use Multiple Runs**: Average results across random seeds |
| 6. **Test Thoroughly**: Use held-out test set for final evaluation |
|
|
| ## 📚 Complete Example |
|
|
| Here's a full workflow from BYOL training to classification: |
|
|
| ```bash |
| # 1. Train BYOL (this takes 4-5 hours on A100) |
| python train_byol_mammo.py |
| |
| # 2. Prepare classification data (create CSVs with labels) |
| # ... prepare train_labels.csv and val_labels.csv ... |
| |
| # 3. Fine-tune for classification (1-2 hours) |
| python train_classification.py \ |
| --byol_checkpoint ./mammogram_byol_best.pth \ |
| --train_csv ./train_labels.csv \ |
| --val_csv ./val_labels.csv \ |
| --tiles_dir ./tiles \ |
| --class_names mass calcification architectural_distortion asymmetry normal \ |
| --output_dir ./classification_results |
| |
| # 4. Run inference on new images |
| python inference_classification.py \ |
| --model_path ./classification_results/best_classification_model.pth \ |
| --images_dir ./new_patient_tiles \ |
| --output_json ./patient_predictions.json |
| ``` |
|
|
| This gives you a complete pipeline from self-supervised pre-training to production-ready classification! 🚀 |
