Exploring the Interplay of Dataset Size and Imbalance on CNN Performance in Healthcare: Using X-rays to Identify COVID-19 Patients

Abstract

Introduction: Convolutional Neural Network (CNN) systems in healthcare are influenced by unbalanced datasets and varying sizes. This article delves into the impact of dataset size, class imbalance, and their interplay on CNN systems, focusing on the size of the training set versus imbalance-a unique perspective compared to the prevailing literature. Furthermore, it addresses scenarios with more than two classification groups, often overlooked but prevalent in practical settings.

Methods: Initially, a CNN was developed to classify lung diseases using X-ray images, distinguishing between healthy individuals and COVID-19 patients. Later, the model was expanded to include pneumonia patients. To evaluate performance, numerous experiments were conducted with varied data sizes and imbalance ratios for both binary and ternary classifications, measuring various indices to validate the model's efficacy.

Results: The study revealed that increasing dataset size positively impacts CNN performance, but this improvement saturates beyond a certain size. A novel finding is that the data balance ratio influences performance more significantly than dataset size. The behavior of three-class classification mirrored that of binary classification, underscoring the importance of balanced datasets for accurate classification.

Conclusions: This study emphasizes the fact that achieving balanced representation in datasets is crucial for optimal CNN performance in healthcare, challenging the conventional focus on dataset size. Balanced datasets improve classification accuracy, both in two-class and three-class scenarios, highlighting the need for data-balancing techniques to improve model reliability and effectiveness.

Motivation: Our study is motivated by a scenario with 100 patient samples, offering two options: a balanced dataset with 200 samples and an unbalanced dataset with 500 samples (400 healthy individuals). We aim to provide insights into the optimal choice based on the interplay between dataset size and imbalance, enriching the discourse for stakeholders interested in achieving optimal model performance.

Limitations: Recognizing a single model's generalizability limitations, we assert that further studies on diverse datasets are needed.

Keywords: CNN; COVID-19; classification biases; imbalanced data; pulmonary disease.

Figure 1

Chest X-ray figures used for the training. (A). (Normal fig.) A normal chest X-ray of a healthy individual. (B). (Lung-opacity fig.) A patient with unilateral right-sided opacity from another cause. (C). (COVID-19 fig.) A COVID-19 patient with bilateral opacities consistent with COVID-19 pneumonia.

Figure 2

CNN architecture—Figure illustrating the detailed architecture of the CNN, as described above.

Figure 3

Consistent patterns across all four data-split methods utilized. As the training set size progressively increases, there is a corresponding improvement in the level of accuracy.

Figure 4

An increase in the IR led to decreased accuracy.

Figure 5

At an IR of 1:99, although the model failed to classify any instances from the minority class correctly, the overall accuracy reached 98.9%.

Figure 6

As the size of the training set increases, the point at which the model’s sensitivity drops below 50% occurs at more extreme imbalance ratios. The black line on the graph represents the sensitivity threshold dropping below 50%. For a 1000-image training set, this occurs around an IR of 20:80. With a 5000-image training set, the threshold is slightly before an IR of 10:90. For a 10,000-image training set, the threshold is just before an IR of 5:95.

Figure 7

As the percentage of the data of a class decreased, the model’s ability to classify that class also decreased. In the balanced dataset, the count of accurately classified images was comparably consistent across each class. However, as the class imbalance intensified, the model tended to classify more images into the majority class. With an extreme imbalance ratio of 1:1:98, the model demonstrated precise classification of the majority class while misclassifying the minority classes.