Medical Image Classification Using Transfer Learning and Network Pruning Algorithms

Deep neural networks show great advancement in recent decades in classifying medical images (such as CT-scans) with high precision to aid disease diagnosis. However, the training of deep neural networks requires significant sample sizes for learning enriched discriminative spatial features. Building a high quality dataset large enough to satisfy model training requirement is a challenging task due to limited disease sample cases, and various data privacy constraints. Therefore in this research, we perform medical image classification using transfer learning based on several well-known deep networks, i.e. GoogLeNet, Resnet and EfficientNet. To tackle data sparsity issues, a Wasserstein Generative Adversarial Network (WGAN) is used to generate new medical image samples to increase the numbers of training instances of the minority classes. The transfer learning process itself also allows the building of strong classifiers by transferring knowledge from the pre-trained image domain to a new medical domain using a small sample size. Moreover, the lottery ticket hypothesis is also used to prune each transfer learning network trained using the new target image data sets. Specifically, the L1 norm unstructured pruning technique is used for network reduction. Hyper-parameter fine-tuning is also performed to identify optimal settings of key network hyper-parameters such as learning rate, batch size and weight decay. A total of 20 trials are used for optimal hyper-parameter selection. Evaluated using multi-class lung X-ray images for pneumonia conditions and brain tumor CT-scans, the fine-tuned EfficientNet model obtains the best brain tumor classification accuracy rate of 96% and a fine-tuned GoogLeNet model with pruning has the highest pneumonia classification accuracy rate of 81.5%.