Multi-CNNs with variational information bottleneck for chest X-ray classification
摘要
Convolutional neural networks (CNNs) have received much attention in the field of medical image classification due to their excellent feature extraction capability, parameter sharing mechanism and spatial structure preservation. Variational information bottleneck (VIB) has the advantages of feature compression, information preservation and prevention of model overfitting. Theoretically, embedding a VIB module into CNNs helps to improve the classification performance of the models. However, traditional methods usually embed VIB as an intermediate feature layer in neural networks. This can result in excessive feature compression, which can filter out information that is critical for classification tasks. Therefore, we propose embedding VIB in several CNNs to provide classification models for chest X-ray images. Specifically, VIB is embedded after the classification layers of DenseNet121, MobileNetV2 and ResNet50 to enhance the performance of medical image classification. First, three publicly available datasets are pre-processed using the contrast limited adaptive histogram equalization (CLAHE), which in turn spatially rotated and horizontal mirroring these datasets in two dimensions to increase the volume of data. Comparisons with other related models for classification of chest X-ray image are then performed. Finally, ablation experiments are performed with and without VIB, as well as with VIB embedded in the interlayer. The above experimental results show that the combination of CNNs with VIB, for binary and multi-classification tasks, achieves satisfactory results in accuracy, precision, recall and F1-score, respectively, which implies that our proposed model has a great potential for application in computer-aided pneumonia diagnosis.