Improving COVID-19 CT classification of CNNs by learning parameter-efficient representation

Abstract

The COVID-19 pandemic continues to spread rapidly over the world and causes a tremendous crisis in global human health and the economy. Its early detection and diagnosis are crucial for controlling the further spread. Many deep learning-based methods have been proposed to assist clinicians in automatic COVID-19 diagnosis based on computed tomography imaging. However, challenges still remain, including low data diversity in existing datasets, and unsatisfied detection resulting from insufficient accuracy and sensitivity of deep learning models. To enhance the data diversity, we design augmentation techniques of incremental levels and apply them to the largest open-access benchmark dataset, COVIDx CT-2A. Meanwhile, similarity regularization (SR) derived from contrastive learning is proposed in this study to enable CNNs to learn more parameter-efficient representations, thus improve the accuracy and sensitivity of CNNs. The results on seven commonly used CNNs demonstrate that CNN performance can be improved stably through applying the designed augmentation and SR techniques. In particular, DenseNet121 with SR achieves an average test accuracy of 99.44% in three trials for three-category classification, including normal, non-COVID-19 pneumonia, and COVID-19 pneumonia. The achieved precision, sensitivity, and specificity for the COVID-19 pneumonia category are 98.40%, 99.59%, and 99.50%, respectively. These statistics suggest that our method has surpassed the existing state-of-the-art methods on the COVIDx CT-2A dataset. Source code is available at https://github.com/YujiaKCL/COVID-CT-Similarity-Regularization.

1. Introduction

The Coronavirus Disease 2019 (COVID-19) has become a worldwide pandemic and infected over $493$ million people till April 2022 [1]. Its increasingly high infectivity and fatality rate due to strain variation are threatening human health and damaging the global economy [2], [3], [4]. The efficient reproductive number of the virus in many countries remains high, as reported in [5], indicating COVID-19 continues spreading quickly around the world. Thereby, a timely and efficient diagnosis is crucial for the treatment of COVID-19 positive patients and the control of further disease spread.

In the early diagnosis of COVID-19 infection, real-time reverse transcription polymerase chain reaction (RT-PCR) is the primary choice due to its convenience and high specificity. However, research results [6], [7], [8] have suggested that RT-PCR is not sensitive enough that some infected patients turned out to be positive even after several negative tests. These false-negative cases might continue to infect their close contacts without isolation or develop into severe illness. Chest computed tomography (CT) has the potential to identify pathological lung changes and thus can serve as a supplementary screening tool to RT-PCR in infection detection, as indicated by [9], [10], [11].

Since the pandemic started, researchers have been exploring the potential of convolutional neural networks (CNNs) in COVID-19 CT classification and reported high accuracy without clinician intervention. CNNs are a kind of deep learning technique dominating computer vision tasks. For example, Gunraj et al. [12] introduced a large-scale open-access COVID-19 CT dataset (COVIDx CT-1) and trained a COVID-19-specific tailored CNN. Panwar et al. [13] utilized transfer learning to inherit cross-domain knowledge to improve the model performance. All these researches reveal that CNNs have the potential to serve as an assistant to help clinicians in COVID CT diagnosis.

Although CNNs have achieved remarkable results in CT diagnosis, challenges remain before they can be put into practical use. Deep learning methods often require large-scale standard datasets, while the existing COVID-19 CT datasets are insufficient. Also, CT scans collected from different institutes have inconsistent characteristics like orientation, brightness, etc. The trained models might be more sensitive to these irrelevant information than the pneumonic pathologies that really matter. Furthermore, the increasingly great capability of CNN-based models may not be fully fulfilled when learning from the limited data sources. Hence, methods for learning more parameter-efficient representations are crucial for mitigating the data insufficiency issue and improving the classification performance.

By addressing these problems above, a more reliable COVID-19 CT classification system can reduce the workload of clinicians and provide more accurate and sensitive computer-aided diagnoses. Motivated by these factors, this study aims to use deep learning techniques to improve the COVID-19 CT classification performance of commonly used CNNs. Particularly, to alleviate the data insufficiency and enhance the data diversity, we design and apply augmentation of incremental levels on the currently largest COVID-19 CT benchmark dataset (COVIDx CT-2A) [14]. Meanwhile, to find the optimal selection of CNN architectures and augmentation combinations, we explore seven commonly used CNN architectures under seven augmentation settings. The CNNs include SqueezeNet1.1 [15], MobileNetV2 [16], DenseNet121 [17], ResNet-18/34/50 [18], and InceptionV3 [19]. Meanwhile, contrastive learning is one promising self-supervised method for enabling deep learning models to learn more parameter-efficient features. We propose the similarity regularization (SR) derived from contrastive learning to learn more parameter-efficient representations and improve CNN classification. The experimental results demonstrate that SR can improve the classification performance of CNNs stably and surpass conventional contrastive learning. Our main contributions are summarized as follows:

• (a)We investigate the impacts of augmentation and model selection in COVID-19 CT classification for three classes, including normal, non-COVID-19 pneumonia (NCP), and COVID-19 pneumonia (CP).
• (b)We propose SR as a regularization term for learning more parameter-efficient representations. Comparisons between seven models with or without SR are conducted. The experimental results demonstrate that SR can improve classification stably without extra introduced model parameters during the test interface.
• (c)Our proposed model, DenseNet121-SR, achieves 99.44% test accuracy, 98.40% precision, 99.59% sensitivity and 99.50% specificity for COVID-19 positive class, achieving the state-of-the-art.
• (d)On other COVID-19 CT datasets, i.e., SARS-CoV2 and COVIDx CT-1, our DenseNet121-SR outperforms the existing methods in terms of efficiency and accuracy.
• (e)We extend the study to seven classic natural datasets and find that our DenseNet121-SR is superior to the original DenseNet121 for all tasks, indicating that our method can be generalized to general classification problems.

The rest of this paper is organized as follows. In Section 2, we review the related works and analyze their pros and cons. Section 3 describes our proposed method. We demonstrate and analyze the experimental results in Section 4 and conduct ablation studies in Section 5. Section 6 discusses our achievements, limitations, and future works. Section 7 draws the conclusion.

2. Related works

2.1. COVID-19 related researches

CNNs are increasingly improving the COVID-19 CT classification with advanced algorithms and enhanced datasets. Numerous CNN-based methods achieving high accuracy have been proposed, indicating the potential of CNNs in assisting practical diagnosis. Some representative methods on four benchmark datasets are listed in Table 1.

Table 1.

Comparison of classification metrics between multiple deep learning methods in four datasets. The precision, sensitivity, and specificity metrics are for COVID-19 positive class only. (The decimal places are kept consistent as reported in the publications.).

Dataset	Method	Params. (M)	Accuracy (%)	Precision (%)	Sensitivity (%)	Specificity (%)
SARS-CoV-2 [20]	Alshazly et al. [21]	86.74	99.4	99.6	99.8	99.6
	Silva et al. [22]	4.78	98.99	99.20	98.80	–
	Kundu et al. [23]	132.86	98.93	98.93	98.93	98.93
	Jaiswal et al. [24]	–	96.25	96.29	96.29	96.21
	Panwar et al. [13]	20.55	94.04	95.30	94.04	95.86
	Jangam et al. [25]	202.87	93.5	89.91	98	–
	Wang et al. [26]	–	90.83	95.75	85.89	–
COVID-CT [27]	Chen et al. [28]	–	88.5	89.9	88.6	–
	He et al. [29]	0.55	86	–	–	–
	Polsinelli et al. [30]	1.26	85.03	85.01	87.55	81.95
	Jangam et al. [25]	202.87	84.73	78.15	94.9	–
	Wang et al. [26]	–	78.69	78.02	79.71	–
COVIDx CT-1 [12]	Gunraj et al. [12]	1.40	99.1	99.7	97.3	99.9
	Ter-Sarkisov [31]	34.14	91.66	90.80	94.75	–
COVIDx CT-2A [14]	Zhao et al. [32]	23.51	99.2	98.5	98.7	99.5
	Gunraj et al. [14]	0.45	98.1	97.2	98.2	98.8
	Gunraj et al. [12]	1.40	94.5	90.2	99.0	95.7

In the COVID-19 CT classification, there exist no gold standard datasets so far. The four widely employed open-access datasets [12], [14], [20], [27] in Table 1 differ in many aspects, including patient/scan distribution, collection sources, dataset size, class numbers, labeling quality, etc. Particularly, COVID-CT [27] and SARS-CoV-2 [20] are two small binary-classification datasets containing $812$ and 2,482 CT scans for COVID-19 positive and non-COVID classes, respectively. Gunraj et al. released a larger dataset COVIDx CT-1 [12] consisting of 104,009 scans for normal, NCP and CP classes upon which the authors later built COVIDx CT-2 [14]. COVIDx CT-2 is the largest existing dataset containing 194,922 CT scans, combined from multiple data sources. Generally, data-driven methods like CNNs depend heavily on dataset size. This can be drawn from the classification metrics in Table 1 that the methods trained on larger datasets can roughly achieve higher performance. To ensure both the data diversity and satisfactory results, our study employs COVIDx CT-2A [14] as our target dataset.

Despite the various datasets, many CNN-based methods have been developed to continuously boost classification performance. In particular, researchers often use transfer learning [13], [21], [22], [23], [24], [25], [33], [34], [35] or ensemble learning [22], [23], [25], [33], [36], [37] to overcome data insufficiency in small-scale datasets like SARS-CoV-2 and COVID-CT. For example, Jaiswal et al. [24] and Panwar et al. [13] utilized transfer learning to pre-train the weights of VGG19 and DenseNet201 on ImageNet and then fine-tuned on SARS-CoV-2, achieving 96.25% and 94.04% accuracy, respectively. Besides, ensemble learning, merging the decisions from multiple models into a more balanced decision, has been widely applied in some works through different merging approaches like weighted sum [25], voting [22], and fuzzy rank-based fusion [23]. However, ensemble learning is rarely applied in large-scale datasets like COVIDx CT-1/2 [12], [14]. Specifically, COVID-Net CT-1/2 L [12], [14] are two light-tailored CNNs whose architectures are finely designed by automatic neural architecture searching. The two models are extremely parameter-efficient and achieved 94.5% and 98.1% accuracy with only 0.45 MB and 1.40 MB parameters, respectively. Another research [32] employed ResNet-50x1 pre-trained on ImageNet-21k and fine-tuned on COVIDx CT-2A, achieving 99.2% accuracy.

Drawn from the reviewed research works above, deep learning models can achieve higher performance in COVID-19 CT classification through the approaches that: (1) train models over data of higher diversity; (2) with finely designed neural networks; (3) ensemble the decisions from multiple models; (4) inherit out-of-domain classification knowledge. Although models can benefit from these aspects, the expensive computational cost of neural architecture searching and large-scale pre-training, and long execution time caused by over-parameterization should be considered as well.

2.2. Contrastive learning

In recent years, supervised deep learning models of increasing complexity and depth have shown great progress in many large-scale applications like ImageNet classification [18], [19]. However, directly applying these models to COVID-19 datasets of smaller scales might cause over-parameterization. It means that model capacities cannot be fully fulfilled, and the extracted representations are not parameter-efficient. One promising approach to addressing the issue is contrastive learning.

In the deep learning field, it is widely recognized that the model performance depends on the quality of their learned representation. Contrastive learning, also known as contrastive self-supervised representation learning, is one framework aiming at learning efficient representations without human-specified labels. In general, the main idea of contrastive learning is to project inputs into an embedding space where the embedded vectors of similar samples are closer while dissimilar ones are apart. More formally, for visual tasks, a pair of views augmented from one image is considered a positive pair, while pairs of views from different images are considered negative pairs. Hence, contrastive learning models aim to maximize the representation similarity between positive pairs and minimize that between negative pairs. In practical tasks, contrastive learning often pre-trains the front representation extractors of deep learning models in a self-supervised manner, and then fine-tunes the pre-trained weights in a conventional supervised manner.

The state-of-the-art contrastive learning frameworks include MoCo  [38], [39], SimCLR [40], [41], SimSiam [42], SwAV [43], BYOL [44], etc. These frameworks mainly differ in terms of the loss function, representation projection, and negative pair formation [42]. And the differences further determine their requirements on the complexity of augmentation policies and batch size. Normally, in order to obtain a satisfactory result, contrastive methods depend on a large batch size to cover enough negative pairs [38], [39], [40], [41]. Among these models, BYOL, SwAV and SimSiam are the contrastive frameworks requiring no negative pairs. In ImageNet linear classification experiments [42], BYOL achieves relatively better performance. This explains that we select BYOL as the basic framework for SR calculation as in Section 3.2.

The success of contrastive learning has emerged some applications in COVID-19 CT diagnosis [28], [29], [45]. He et al. [29] employed a MoCo-like [38] framework to enhance the CT scan representations extracted by DenseNet169 and fine-tuned the network, achieving 86% accuracy on COVID-CT [27]. Similarly, Chen et al. employed the MoCo-v2-like [39] framework on the same dataset and reached 88.5% accuracy within six shots. Li et al. [45] put the contrastive loss as a regularization term and trained their CMT-CNN in an end-to-end manner, obtaining 93.46% accuracy. These studies suggest contrastive learning can boost classification performance by learning more efficient representations.

3. Method

3.1. Augmentation of incremental levels

Data augmentation is vital for improving the performance of deep learning models, especially for contrastive learning [39], [40]. However, the optimal selection for COVID-19 CT augmentation has not been studied. Inspired by the literature in Section 2, we design and evaluate a series of augmentation operations of incremental levels as follows where “+” denotes the appended augmentation based on the previous level:

Level 0

No augmentation.

Level 1

+ RandomResizedCrop: Randomly obtain an image crop of size in the range $[0.08,1]$ of the original size 256 × 256, and then randomly scale the crop according to an aspect ratio in the range $[3/4,4/3]$. The scaled crop is finally resized to the original size.

Level 2

+ Horizontal Flip. Randomly flip the input image horizontally with 50% using probability.

Level 3

+ RandAugment [46]: Randomly apply rand augment twice with magnitude $9$ and magnitude standard deviation 0.5.

Level 4

+ Random Erasing [47]: Select a rectangle region of the input image and do pixel-wise erasing with 25% using probability. The size of the selected region are randomly picked in the range $[0.02,1/3]$ of the image size.

Level 5

+ Mixup [48]: Mix two in-batch images up with a ratio $\lambda$ subjecting to a beta distribution, $\lambda \sim B(1,1)$. The mixup process for images $I_A$ and $I_B$ can be formatted as $I_A(x,y)=\lambda I_A(x,y)+(1-\lambda )I_B(x,y)$, where $(x,y)$ denotes the pixel coordinate.

Level 6

+ CutMix [49]: Switch from Mixup to Cutmix with 50% probability. Randomly replace a square region in the original image with a region in another in-batch image. The region size is randomly determined, subject to the squared root of a beta distribution $B(1,1)$.

The visualization of the augmented scans is demonstrated in Fig. 2. Specifically, RandomResizedCrop and horizontal flip are two commonly used augmentation operations in both supervised [18], [19] and self-supervised learning [39], [40], [41], [44]. Since contrastive learning requires more complicated augmentation [40], the two stronger augmentations, RandAugment and Random Erasing, are further introduced in levels 3 and 4. Their implementations and parameters refer to [50], [51]. In levels 5 and 6, mixup and cutmix are two augmentations enabling higher data diversity by fusing in-batch images. In these two levels, we mainly experiment on whether such sample-fusing augmentations can improve COVID-19 classification. By comparing the performance of models under these incremental augmentation levels, an appropriate augmentation strategy for COVID-19 CT scans can be established.

Fig. 2.

Illustration of applied incremental augmentation of six levels. On the right side, six groups of images are augmented from the same left COVID-19 positive scan. These groups from top to down are in the augmentation levels from 1 to 6. The bottom left normal scan is the auxiliary original image that only participated in Mixup/Cutmix augmentation in levels 5 and 6. The left two scans are from COVIDx CT-2A.

3.2. Similarity regularization

Most mainstream conventional CNNs contain two parts, a representation extractor $f$ and a followed fully connected layer $\text{FC}$. The extractor $f$ aims to extract the distinguishable representations of given inputs, and $\text{FC}$ predicts the class probability distribution by summarizing the extracted representations. This forward propagation is demonstrated as the top branch in Fig. 1. More formally, the input image $x$ is first transformed to a view $v$ by a random on-the-fly augmentation operation $t\sim T$ where $T$ denotes an infinite collection of augmentation operations. Subsequently, the representation extractor $f$ converts the input view $v$ to a representation embedding vector $h=f\left(v\right)$. $\text{FC}$ predicts the class probability distribution based on its obtained representation, $\hat{y}=\text{FC}\left(h\right)$. The training target of such a classifier is to minimize the class probability distribution distance between the prediction $\hat{y}$ and the ground truth $y$ according to the cross-entropy loss in Eq. (1), where $i\in \{0,1,2\}$ denotes the class index.

$$\mathcal{H}(y,\hat{y})=-\sum _{i=0}^2{y}_{i}log{\hat{y}}_i$$ (1)

Fig. 1.

The overall structure of the models with our proposed similarity regularization in training interface. The two projectors, $g_1$ and $g_2$, and the online predictor $p$ are implemented by non-linear MLPs. After training, only the online encoder $f_1$ and fully connected layer $\text{FC}$ are preserved in testing.

In this conventional fully supervised scenario, the trained representations aim at better projecting to human-specified class distribution. However, this manner affects data efficiency, robustness or generalization [52]. Instead, contrastive learning enables learning more parameter-efficient representations from inputs themselves instead of the specified annotations. We thus incorporate it in common CNNs to improve their representation learning ability.

The overall structure of our method is illustrated in Fig. 1. We keep the conventional supervised classifier unchanged in the top branch while introducing a contrastive learning framework in the bottom branch. As in Section 2.2, contrastive learning is to maximize the representation similarity between positive pairs. We punish the positive-pair representation distance as a regularization term beside the cross-entropy loss, naming the term similarity regularization (SR).

Particularly, the contrastive framework is a siamese network like most mainstream frameworks [40], [41], [42], [44], consisting of an online network and a target network. The target network can be seen as a moving average of the online one. Given two views $v_1$ and $v_2$ augmented from the same input image $x$, the representation extractors $f_1$ and $f_2$ in two networks extract their corresponding latent representation vectors, $h_1=f_1\left(v_1\right)$ and $h_2=f_2\left(v_2\right)$. To avoid representations heavily affected by SR, the representation vectors then projected to another embedding space where $z_1=g_1\left(h_1\right)$ and $z_2=g_2\left(h_2\right)$, as in [41], [44]. Since the projectors $g_1$ and $g_2$ share slightly different feature spaces, the online projection $z_1$ is further projected to $p\left(z_1\right)$ of same dimension via online predictor $p$. The cosine representation similarity $\mathcal{S}$ of value in range $[-1,1]$ can be measured according to Eq. (2).

$$\mathcal{S}(p\left(z_1\right),z_2)=\frac{〈p\left(z_1\right),z_2〉}{\Vert p\left(z_1\right){\Vert }_2\cdot \Vert z_2{\Vert }_2}$$ (2)

where $〈\cdot ,\cdot 〉$ and $\Vert \cdot {\Vert }_2$ are inner product and $L^2$ norm notations, respectively. A higher value indicates two vectors are of higher similarity. To penalize a low cosine similarity between positive pairs and scale the penalty in range $[0,1]$, SR can be calculated as Eq. (3).

$$\mathcal{D}(p\left(z_1\right),z_2)=2-2\mathcal{S}(p\left(z_1\right),z_2)=2-\frac{2〈p\left(z_1\right),z_2〉}{\Vert p\left(z_1\right){\Vert }_2\cdot \Vert z_2{\Vert }_2}$$ (3)

Hence, for a positive pair $(v_1,v_2)$, its total loss containing both cross-entropy loss and SR is written as in Eq. (4).

$$\mathcal{L}(z_1,z_2,y,\hat{y})=(1-\gamma )\mathcal{H}(y,\hat{y})+\gamma \mathcal{D}(p\left(z_1\right),z_2)$$ (4)

where $\gamma$ is a scale factor for balancing the conventional cross-entropy loss and the introduced SR.

$(v_2,v_1)$ is the symmetric positive pair with respect to $(v_1,v_2)$. We calculate the losses for both symmetric pairs and take their mean as the final loss for fast convergence.

SR as a regularization term may raise concern if it will dominate the combined loss and thus degrade the classification. To remove the concern and find an appropriate scheduler for $\gamma$, we design three strategies as listed below. $i$ denotes the current training iteration number.

Constant (default)

: $\gamma$ is set to be a constant value during all training iterations, 0.5 by default.

Linear Decay

: $\gamma$ decays linearly to a minimum value ${\gamma }_{min}=0.01$ along $N$ training iterations according to Eq. (5).

$${\gamma }_i={\gamma }_{min}+(1-\frac{i}{N})(1-{\gamma }_{min})$$ (5)

Cosine Decay

: $\gamma$ decays to a minimum value ${\gamma }_{min}=0.01$ along $N$ training iterations according to cosine annealing scheduler as in Eq. (6).

$${\gamma }_i={\gamma }_{min}+\frac{1}{2}(1+cos\frac{i\pi }{N})(1-{\gamma }_{min})$$ (6)

Besides, after training, we throw away all the components except the online representation extractor $f_1$ and the fully connected layer $\text{FC}$. Hence, introducing SR in training will not slow down the test interface. The training pseudocode of models with SR is demonstrated in Algorithm 1.

4. Results and analysis

4.1. Dataset description

In this paper, we mainly train and evaluate our proposed method using the largest existing open-access COVID-19 CT dataset, COVIDx CT-2A.1 Specifically, the dataset contains three classes, including normal, non-COVID-19 pneumonia (NCP), and COVID-19 pneumonia (CP). Its class distribution is summarized in Table 2. The dataset is of high diversity, containing scans of 3,745 patients from eight open-access sources. It should be noted that the scans from the same patient are in one subset, preventing information leakage from training to validation or testing.

Table 2.

Class distribution of the employed COVIDx CT-2A dataset.

Set	Normal	NCP	CP	Total
Training	35,996	25,496	82,286	143,778
Validation	11,842	7400	6244	25,486
Testing	12,245	7395	6018	25,658

4.2. Experimental setting

In this paper, we keep the hyper-parameters consistent in all experiments for fair comparisons. The codes are implemented by PyTorch. We implement the CNN backbones and image augmentation by torchvision and timm [51] libraries, respectively. For acceleration, we train models on Torch distributed data parallelism on four Nvidia V100 GPUs with apex mixed precision of level $O1$. Besides, to alleviate the randomness concern, we obtain the experimental statistics by averaging the measurements in three distinct trials.

During training, CT scans are resized to 256 × 256 in $3$ channels using bicubic interpolation and normalized by ImageNet mean and std. In the test interface, 256 × 256 CT scans are cropped from the center of resized 293 × 293 original images. This is empirically good as the center crop can preserve the main lung regions. To avoid models from being too confident in one-class prediction, label smoothing [19] of smoothing factor 0.1 is applied in the cross-entropy in augmentation levels $0-4$. While in augmentation level $5$ or $6$, in-batch paired labels are mixed up based on mixed inputs (See [48], [49] for more details).

The optimizer we used is Adam with 10⁻⁶ weight decay. After a 5-epoch linear warmup [53] from $5\times 10^{-7}$, we use cosine annealing scheduler to decay the learning rate from $5\times 10^{-4}$ to $5\times 10^{-7}$ in the later $45$ epochs. The batch size is set to $64$ in each process. Besides, the gradients are clipped to be no larger than 5.0 to avoid overflow.

In the SR calculation, the projectors $g_1,g_2$ and predictor $p$ have the same multi-layer perceptron (MLP) architecture that consists of two linear layers connected by a batch normalization layer and a ReLU activation layer. The front linear layer projects the inputs to $512$-D embedding vectors and the later linear layer outputs $128$-D vectors. The analysis for the dimension setting is in Section 5.3. The momentum rate $\beta$ for updating $f_2$ and $g_2$ is 0.99, a median value among contrastive frameworks [38], [42], [43], [44].

4.3. Results of ResNets under incremental augmentation levels

We first compare the performance of ResNets with or without SR under the incremental augmentation levels designed in Section 3.1 to determine an appropriate augmentation policy for the coming experiments. The averaged test accuracies are listed in Table 3. Since SR requires calculating the similarity between two augmented views, models with SR cannot be implemented under augmentation level $0$.

Table 3.

Test accuracy comparison between original ResNets and ResNets with proposed SR under incremental augmentation levels. ResNet is abbreviated as R. The scale factor scheduler for scaling SR is the default constant scheduler $\gamma =0.5$.

Augmentation	Method	CNN backbone
		R18	R34	R50
Level 0	Original	91.28	92.28	91.57
Level 1	Original	99.10	99.00	98.89
	＋SR (Ours)	99.23	99.27	99.09
Level 2	Original	99.17	99.11	99.00
	＋SR (Ours)	99.27	99.39	99.19
Level 3	Original	99.11	99.06	99.08
	＋SR (Ours)	99.26	99.29	99.20
Level 4	Original	99.12	99.09	99.13
	＋SR (Ours)	$\mathbf{99.40}$	$\mathbf{99.43}$	$\mathbf{99.31}$
Level 5	Original	97.49	97.34	97.66
Level 6	Original	98.50	98.69	98.91

Table 3 shows that the original ResNets achieve the highest accuracy in level $2$ and cannot be improved in the following levels, heavily degraded in levels $5$ and $6$. The degradation may result from the fact that sample-fusing augmentation sometimes transfers the pneumonic pathologies from CP/NCP cases to normal cases. We thus do not perform SR in levels $5$ and $6$. Different from the original ResNets, ResNets with SR continue to improve after level $2$ and achieve the highest accuracy in level $4$. This is consistent with the findings in many contrastive learning research works that contrastive learning requires stronger augmentation than supervised models [39], [40], [54]. Hence, we select level $4$ as the basic augmentation level for the following experiments. Overall, it is observed that SR can improve the classification performance of ResNets stably under all augmentation levels from $1$ to $4$.

4.4. Results under augmentation level 4

In this section, we extend the experiments to seven widely used CNNs, including SqueezeNet1.1 [15], MobileNetV2 [16], DenseNet121  [17], ResNet-18/34/50 [18], and InceptionV3 [19]. The experiments are under augmentation level 4 and a constant scale factor scheduler $\gamma =0.5$.

From the results in Table 4, it can be seen that all our models with SR surpass the original models in terms of averaged test accuracy. The best model, DenseNet121 with SR, achieves 99.44% accuracy with $7.33M$ parameters. Note that the extra parameters will be thrown away after training so that the parameters in the test interface are consistent for a model with or without SR. Another observation is that, in COVID-19 CT classification, the model performance is not strictly proportional to its capacity despite model architecture. This suggests that the fine design of model architecture rather than simply expanding depth or width is more valuable in this task, as supported in [12], [14], [30].

Table 4.

Test accuracy of seven CNNs with or without SR under augmentation level 4, ordered by the number of parameters. The scale factor scheduler for scaling SR is the default constant strategy $\gamma =0.5$. Note that the extra parameters introduced by SR will be discarded in test interface after training.

Method	Metric	CNN backbone
		SqueezeNet	MobileNet	DenseNet	ResNet18	ResNet34	InceptionV3	ResNet50
Original	Params. (M)	0.69	2.12	6.63	10.66	20.30	20.78	22.42
	Acc. (%)	98.22	99.02	99.05	99.08	99.06	99.22	99.12
＋SR (Ours)	Params. (M)	1.13	2.94	7.33	11.10	20.74	21.97	23.62
	Acc. (%)	98.41	99.18	$\mathbf{99.44}$	99.39	99.43	99.32	99.31

Fig. 3 shows the confusion matrices for DenseNet121-SR in three training trials. Based on the matrices, we measure the performance of the model in terms of averaged accuracy, precision, sensitivity, and specificity, as listed in Table 5. The results show that our DenseNet121-SR has outperformed the state-of-the-art models in nearly all measurements. Specifically, DenseNet121-SR achieves a high sensitivity 99.59% for COVID-19 positive class, indicating that the model has the potential to efficiently avoid COVID-19 positive patients from being wrongly diagnosed.

Fig. 3.

Confusion matrix for DenseNet121-SR in three trials.

Table 5.

Comparison of DenseNet121-SR with the state-of-the-art methods on COVIDx CT-2A dataset.

Method	Accuracy (%)	Precision (%)			Sensitivity (%)			Specificity (%)
		Normal	NCP	CP	Normal	NCP	CP	Normal	NCP	CP
COVID-Net CT-1 [12]	94.5	96.1	97.6	90.2	98.8	80.2	99.0	96.3	99.4	95.7
COVID-Net CT-2 S [14]	97.9	99.3	96.4	97.0	98.9	95.7	98.1	99.3	98.9	98.8
COVID-Net CT-2 L [14]	98.1	99.4	96.7	97.2	99.0	96.2	98.2	99.5	99.0	98.8
Bit-M [32]	99.2	99.8	98.9	$\mathbf{98.5}$	$\mathbf{99.3}$	99.6	98.7	99.8	99.6	99.5
DenseNet121-SR (Ours)	$\mathbf{99.44}$	$\mathbf{99.89}$	$\mathbf{99.55}$	98.40	99.12	$\mathbf{99.83}$	$\mathbf{99.59}$	$\mathbf{99.91}$	$\mathbf{99.82}$	$\mathbf{99.50}$

To better understand the behavior of our model, we visualize the attention of DenseNet121-SR on three CT scans in different classes as in Fig. 4. It can be observed that our model mainly focuses its attention on some suspicious regions where the pathologies may exist.

Fig. 4.

Attention of DenseNet121-SR visualized by Grad-CAM. The three groups of CT scans and heatmaps from left to right are in class normal, NCP, and CP, respectively. The highlighted parts are the regions based on which CNNs classify the CT scans.

4.5. Results on other datasets

On other COVID-19 CT datasets Based on the experimental results aforementioned, we extend our method to the other two COVID-19 CT datasets, i.e., SARS-CoV2 and COVIDx CT-1. It should be noted that, for SARS-CoV2, we train DenseNet121-SR over $200$ epochs with weights pre-trained on ImageNet because SARS-CoV2 contains much fewer CT scans than the others. The results as listed in Table 7 show that our method can be generalized to other datasets and can achieve a high classification performance. Comparing to the methods listed in Table 1, our DenseNet121-SR with only 6.63 MB parameters is more parameter-efficient and outperforms the reviewed methods.

Table 7.

Classification results of DenseNet121-SR on SARS-CoV2 and Covidx CT-1 datasets. The precision, sensitivity, and specificity metrics are for COVID-19 positive class only.

Dataset	Acc (%)	Prec (%)	Sens (%)	Spec (%)
SARS-CoV2	99.20	99.47	98.93	99.46
COVIDx CT-1	99.78	99.56	99.84	99.74

On classic natural datasets Besides, extensive experiments are conducted over seven natural datasets to further evaluate the generalization ability of our method. To evaluate the effect of SR fairly, we keep the setting unchanged as in Section 4.2 and initialize the model weights as pre-trained on ImageNet. Table 6 demonstrates the classification accuracy of DenseNet121 with or without SR on the seven datasets, including FGVC Aircraft [55], CIFAR10/100 [56], Describable Textures Dataset (DTD) [57], Oxford 102 Flowers [58], Oxford-IIIT Pets [59], and Stanford Cars [60]. It shows that DenseNet121-SR is superior to the original model in all the tasks, indicating our proposed SR can be generalized to general classification problems.

Table 6.

Classification accuracy of DenseNet121 with or without SR over seven classic natural datasets.

	Aircraft	CIFAR10	CIFAR100	DTD	Flowers102	OxfordIIITPet	StanfordCars
DenseNet121	88.15	94.45	85.08	70.60	93.17	92.47	92.46
DenseNet121-SR (Ours)	$\mathbf{88.18}$	$\mathbf{94.47}$	$\mathbf{85.08}$	$\mathbf{71.01}$	$\mathbf{94.42}$	$\mathbf{92.88}$	$\mathbf{92.55}$

5. Ablation study

The following ablation studies are conducted to better investigate the effects of our proposed SR.

5.1. Ablation to self-supervised learning

Fully self-supervised learning Contrastive learning is widely adopted in pre-training CNNs that are fine-tuned later for downstream tasks. In our methods, we turn the process to an end-to-end manner by regularizing CNNs with proposed SR derived from contrastive learning. Hence, a comparison between SR and conventional contrastive learning is necessary. Specifically, we design and measure the following methods for comparison as follows.

• (a)Linear Evaluation. First pre-train the representation extractor $f$ weights of which are frozen in the later $\text{FC}$ fine-tuning. The pre-training process is equivalent to setting $\gamma =1$ in all training epochs as in Algorithm 1, and then fine-tuning only the linear layers $\text{FC}$ as usual. The fine-tuning hyper-parameters include: $256$ batch size, learning rate decays from $40$ to $4\times 10^{-6}$ according to cosine decay scheduler [53]. The optimizer used is SGD. Linear evaluation is simply conducted for verifying the effects of contrastive learning in this task.
• (b)Two-stage training (self-supervised contrastive learning followed by conventional supervised learning). This way first pre-trains the representation extractor $f$ and then trains the entire CNN with pre-trained weights. The hyper-parameters are consistent with others as in Section 4.2.
• (c)Apply SR to ResNets with a default constant $\gamma =0.5$.

The results for the designs above are listed in Table 8. It can be observed that contrastive learning can learn efficient representations that even a simple linear evaluation on the pre-trained representation extractor can achieve over 92% test accuracy. For the second method, two-stage contrastive learning, the pre-trained weights from the representation extractor might be hard to maintain in the later training phase. Our introduced SR maintains the representation by explicitly penalizing the representation difference for positive pairs. The results in Table 8 verify that ResNets with SR surpass the two-stage contrastive learning method in most experiments. Besides, it is worth noting that the end-to-end training with SR does not require pre-training and thus saves computational resources.

Table 8.

Test accuracy comparison between two-stage training and end-to-end training with SR under incremental augmentation levels. The scale factor scheduler for scaling SR is the default constant scheduler $\gamma =0.5$.

Augmentation	Method	CNN backbone
		R18	R34	R50
Level 1	Linear Eval	95.47	92.89	92.32
	Two-stage	99.17	99.17	99.04
	＋SR (Ours)	$\mathbf{99.23}$	$\mathbf{99.27}$	$\mathbf{99.09}$
Level 2	Linear Eval	95.76	92.03	94.02
	Two-stage	99.25	99.23	$\mathbf{99.20}$
	＋SR (Ours)	$\mathbf{99.27}$	$\mathbf{99.39}$	99.19
Level 3	Linear Eval	93.92	93.88	95.08
	Two-stage	99.26	99.29	$\mathbf{99.24}$
	＋SR (Ours)	$\mathbf{99.26}$	$\mathbf{99.29}$	99.20
Level 4	Linear Eval	93.12	93.23	94.53
	Two-stage	99.38	99.30	99.18
	＋SR (Ours)	$\mathbf{99.40}$	$\mathbf{99.43}$	$\mathbf{99.31}$

5.2. Ablation to decay strategy for $\gamma$

The two-stage contrastive learning method can be approximated to run the SR Algorithm 1 with $\gamma =1$ in pre-training and $\gamma =0$ in fine-tuning. The sharp fall of $\gamma$ may destroy the maintained representation space obtained in contrastive pre-training. To avoid the potential negative impact, we designed two mild $\gamma$ decay strategies in Section 3.2 despite the constant $\gamma$ strategy. From the results demonstrated in Fig. 5, we can conclude that SR with all designed $\gamma$ strategies can stably improve classification accuracy. And SR is insensitive to the $\gamma$ strategy setting since all strategies have comparable performance. Due to the simplicity of the $\gamma$ strategy ($\gamma =0.5$ in all iterations) and its slight superiority in level $4$ augmentation, we select it as the default strategy in our experiments.

Fig. 5.

Accuracy of ResNets obtained with different $\gamma$ decay strategies under incremental augmentation levels from 1 to 4. The baselines are the original models without introduced SR. The $\gamma$ value for the constant strategy is 0.5 by default while $\gamma$ decays from 1.0 to 0.01 in linear and cosine strategies.

The ablation studies find that a constant strategy, $\gamma =0.5$, can achieve the relatively highest performance among the three strategies under level $4$ augmentation.

The constant $\gamma$ value still requires studies for finding its effects on model performance. We thus vary the $\gamma$ value in constant strategy from 0.1 to 0.9 with 0.2 interval and repeat the experiments for CNNs with SR under augmentation level $4$. As shown in Fig. 6, SR can improve the CNN classification performance when $\gamma$ value is in an appropriate range near $[0.5,0.7]$. In particular, a smaller $\gamma$ cannot fully fulfills the advantage of SR and sometimes even degrades the model capacity as in SqueezeNet1.1 case. Meanwhile, setting $\gamma$ to a large value like 0.9 is also risky since SR dominates the total loss while the primary cross entropy for classification is slighted.

Fig. 6.

Accuracy of seven CNNs with SR controlled by constant $\gamma$ scale factor strategy, under augmentation level $4$. $\gamma$ varies from 0.1 to 0.9 with 0.2 interval.

5.3. Ablation to projection size in SR

The output dimension or named projection size, of both the projector and predictor in SR calculation is set to $128$ as default. We keep the hidden dimension $512$ unchanged to avoid redundant computation while varying the projection size to analyze its effect in terms of classification accuracy. As visualized in Fig. 7, the differences between the classification accuracy for models except for SqueezeNet are small ($\le 0.2\text{\%}$). This indicates that the hyper-parameter setting in our proposed SR is robust.

Fig. 7.

Impact of projection size in SR calculation.

6. Discussion

Since COVID-19 grows rapidly worldwide, designing efficient and accurate classification systems is essential. Although some methods [12], [14], [21], [22], [32] have claimed a high classification accuracy ($\approx$99%) on multiple datasets, we argue that even a slight improvement can mitigate further infection. Meanwhile, some high-performance methods require considerable computational resources, making them hard to be deployed into practical healthcare systems. Hence, designing more efficient models with affordable training parameters should also be noted.

This paper mainly proposes an incremental augmentation strategy and SR to improve the CNN classification performance on three COVID-19 CT datasets. The results illustrate that appropriate augmentation can significantly alleviate the data limitation problem in COVID-19 CT classification. Meanwhile, our proposed SR further improves the classification performance of seven CNNs by enhancing their representation learning ability. Specifically, on the largest dataset COVIDx CT-2A, our model DenseNet121-SR achieves 99.44% accuracy and 99.59% sensitivity with only 6.63 MB parameters in the test interface, outperforming all the reviewed state-of-the-art methods. Overall, the designed augmentation strategy enhances the COVID-19 CT dataset diversity approximately and the proposed SR better fulfills the representation learning ability of CNNs within definite parameters, both leading to the high classification performance of our work.

Besides, we evaluate DenseNet121-SR on the other two datasets, achieving 99.78% and 99.20% accuracy on COVIDx CT-1 and SARS-CoV2, respectively. To further justify the effect of SR, we extend the DenseNet121-SR to seven classic natural datasets, illustrating SR can be generalized to general classification tasks. Furthermore, since SR derives from contrastive learning, we compare traditional contrastive learning and end-to-end training with SR in ablation studies. The comparison demonstrates that SR is superior in classification accuracy and training efficiency and is robust to its hyper-parameter setting.

Despite the achieved promising performance, the limitations of our method exist. Our method requires either large amounts of training data or pre-training on other large-scale datasets. The high performance of our models with SR partly owes to the efforts of workers collecting numerous CT scans. For smaller-scale datasets like SARS-CoV2, the backbone of our method requires pre-training. The pre-training on ImageNet helps improve the accuracy from around 98% to 99.20% in the DenseNet121-SR case. Besides, we can hardly evaluate other contrastive frameworks due to the lack of computational resources. Meanwhile, we cannot redesign the CNN backbones to better balance computational efficiency and classification performance because of the substantial computational loads of neural architecture searching and pre-training.

Before deploying to clinical practice, computer-aided methods must be evaluated regarding their robustness and generalization. Compared to the existing online CT datasets, CT scans in practical scenarios are obtained by more diverse imaging devices under changeable environments. It is significant for the methods to adapt to these unfamiliar scenarios. Also, the interpretability of the methods for inferring the decision principle can make the methods more trustable. Besides, with the evolution of coronavirus, whether the pathological changes can still be observed clearly in CT scans should be concerned as well. Computer-aided methods can convince clinicians and patients as assisting tools in practical uses only if the issues above are addressed.

For future work, we will explore redesigning the network backbone, pre-training the redesigned backbones on large-scale datasets, and making networks more explainable in COVID-19 CT diagnosis.

7. Conclusion

This paper aims to improve the CNN performance for COVID-19 CT classification by enabling CNNs to learn parameter-efficient representations from CT scans. We propose the SR technique derived from contrastive learning and apply it to seven commonly used CNNs. The experimental results show that SR can stably improve the CNN classification performance. Together with a well-designed augmentation strategy, our model DenseNet121-SR with 6.63 MB parameters outperforms the existing methods on three COVID-19 CT datasets, including SARS-CoV2, COVIDx CT-1, and COVIDx CT-2A. Specifically, on the largest available dataset COVIDx CT-2A, DenseNet121-SR achieves 99.44% accuracy, 98.40% precision, 99.59% sensitivity, and 99.50% specificity for the COVID-19 pneumonia category. Furthermore, the extensive experiments on seven classic natural datasets demonstrate that SR can be generalized to common classification problems.

CRediT authorship contribution statement

Yujia Xu: Conceptualization, Methodology, Writing – original draft. Hak-Keung Lam: Method improvement, Writing – review & editing, Supervision. Guangyu Jia: Method improvement. Jian Jiang: Writing – review & editing. Junkai Liao: Writing – review & editing. Xinqi Bao: Writing – review & editing.

Declaration of Competing Interest

The authors declare that they have no known competing financial interests or personal relationships that could have appeared to influence the work reported in this paper.