A New Lightweight Architecture and a Class Imbalance Aware Loss Function for Multi-label Classification of Intracranial Hemorrhages

Abstract

Deep learning algorithms have proven effective in solving many medical imaging tasks in recent years. The design of lightweight neural networks is gaining importance in the medical imaging community as not many hospitals and clinics are equipped with high computational resources to deploy large deep learning algorithms. Also, medical imaging data often comes with high class imbalance and thus there is a high necessity to develop deep learning models that can address this issue. With this motivation, a resource-efficient deep learning model called Lightweight-Fully Convolutional Network (LightFCN) is developed which can be deployed in clinical settings with limited computational resources. Label Distribution Aware Margin loss (LDAM) is used in the context of medical imaging for the first time for multi-label classification with class imbalance. The proposed model has a smaller memory footprint, a smaller number of parameters, lesser inference time and fewer Floating Point Operations (FLOPS) when compared to state-of-the-art models, without compromising on performance and can be deployed in clinical settings with limited computational resources. The model and the performance of the loss function are evaluated on the task of Intracranial Hemorrhage (ICH) classification on CT scans, and the model was deployed on a Raspberry Pi 4B (8 GB), on which inference times were compared. It is found that the proposed model significantly reduced the number of model parameters by a factor of 26, and reduced the inference time by a factor of 3, when compared to the popular lightweight network MobileNetV2.