M.S. AAI Capstone Chronicles 2024

CNN Lung Disease Classification

11

‘ModelCheckpoint’ tool utilized to save the best-performing models based on validation loss.

Transfer Learning with VGG16:

In addition to the previous CNNs, we explored transfer learning using the pre-trained VGG16 model (Simonyan & Zisserman 2015). The VGG16 model, with its weights pre-trained on the ImageNet dataset, was used as a feature extractor. The original classification layers of VGG16 were removed, and a custom classifier head was added, consisting of a flatten layer, a dropout layer, two dense layers with ReLU activation, and a final sigmoid output layer for multi-label classification. This approach leveraged the feature representations learned by VGG16 on a massive dataset, which was particularly useful given the smaller size of our chest X-ray sample. All models were trained using the binary cross-entropy loss function with a threshold of 0.3, which is a common loss function for multi-label classification problems. This lower threshold was chosen to improve recall, which is crucial in medical diagnoses where identifying all possible cases of a condition may be more critical than avoiding false positives (Folio3AI, n.d.). Accuracy, F1 and AUC-ROC scores were used as evaluation metrics, with AUC-ROC scores included to understand model performance across different labels. The training procedure involved monitoring validation loss for early stopping, and experiments with larger image sizes (e.g., 512 x 512 and 300 x 300) showed potential for improved performance, though they required significantly more memory. While grayscale conversion was also explored, it did not

169