Fine-tuned EfficientNet-B0 model that classifies lung CT-scan images into four categories: adenocarcinoma, large.cell.carcinoma, normal, squamous.cell.carcinoma.
Dataset
Trained on the Lung Cancer 4 Types CT Scan Image Dataset, published by Kabil007 on Hugging Face: huggingface.co/datasets/Kabil007/LungCancer4Types. This model repository does not redistribute the dataset itself — only the fine-tuned weights, training code, and evaluation results. Please refer to the original Hugging Face dataset page for the data and its license terms.
Model Description
Architecture:torchvision.models.efficientnet_b0, pretrained on ImageNet, with the final classifier layer replaced to output 4 classes.
Regularization: the EfficientNet-B0 classifier head keeps its built-in Dropout(p=0.2) before the final linear layer, active throughout fine-tuning to reduce overfitting on the 613-image training set.
Input size: 224x224 RGB, normalized with standard ImageNet mean/std.
Augmentation (train only): random horizontal flip, random rotation (±10°).
Training Procedure
Setting
Value
Optimizer
Adam
Learning rate
1e-4
Weight decay
1e-4
Batch size
16
Max epochs
20
Early stopping
patience = 5, monitored on validation macro F1
Checkpoint selection
best validation macro F1
Seed
42
Training stopped early at epoch 11 once validation macro F1 stopped improving. The checkpoint saved corresponds to the best validation macro F1 epoch, not the final epoch — this matters because train loss continued decreasing while validation performance started to plateau, indicating mild overfitting.
Known Limitation: adenocarcinoma vs. squamous cell carcinoma
The confusion matrix shows that the model's errors concentrate mostly between adenocarcinoma and squamous.cell.carcinoma. In particular, 22 adenocarcinoma images were predicted as squamous cell carcinoma, while 2 squamous cell carcinoma images were predicted as adenocarcinoma.
This is the model's main weak spot: squamous.cell.carcinoma has high recall, but its precision is lower because a noticeable number of adenocarcinoma samples are pulled into that class. This suggests that the model learns strong features for detecting squamous cell carcinoma, but still struggles to separate it cleanly from adenocarcinoma.
A related effect appears with large.cell.carcinoma, which has the smallest training set among cancer classes. Its recall is lower than its precision, meaning the model sometimes misses large cell carcinoma cases and predicts them as another cancer subtype.
In short: overall performance is solid, especially for the normal class, but the main limitation is the small dataset size and the visual similarity between cancer subtypes. With only 613 training images, the model has limited examples for learning the finer differences between adenocarcinoma, squamous.cell.carcinoma, and large.cell.carcinoma. More labeled training data, especially for the cancer subtype classes, would likely be the highest-leverage improvement.
This model is trained for research and educational purposes only. It is not a diagnostic tool and has not been validated for clinical use. Predictions should never be used to make or support real medical decisions without review by a qualified pathologist/radiologist.
