Determining Top Fully Connected Layer’s Hidden Neuron Count for Transfer Learning, Using Knowledge Distillation: a Case Study on Chest X-Ray Classification of Pneumonia and COVID-19

Abstract

Deep convolutional neural network (CNN)-assisted classification of images is one of the most discussed topics in recent years. Continuously innovation of neural network architectures is making it more correct and efficient every day. But training a neural network from scratch is very time-consuming and requires a lot of sophisticated computational equipment and power. So, using some pre-trained neural network as feature extractor for any image classification task or “transfer learning” is a very popular approach that saves time and computational power for practical use of CNNs. In this paper, an efficient way of building full model from any pre-trained model with high accuracy and low memory is proposed using knowledge distillation. Using the distilled knowledge of the last layer of pre-trained networks passes through fully connected layers with different hidden layers, followed by Softmax layer. The accuracies of student networks are mildly lesser than the whole models, but accuracy of student models clearly indicates the accuracy of the real network. In this way, the best number of hidden layers for dense layer for that pre-trained network with best accuracy and no-overfitting can be found with less time. Here, VGG16 and VGG19 (pre-trained upon “ImageNet” dataset) is tested upon chest X-rays (pneumonia and COVID-19). For finding the best total number of hidden layers, it saves nearly 44 min for VGG19 and 36 min and 37 s for VGG16 feature extractor.

Introduction

Pneumonia is the single leading cause of childhood mortality. About 2 million children under 5 years old, die due to pneumonia every year around the world according to the World Health Organization (WHO). It kills more children than AIDS, malaria and measles [1]. Almost 95% of childhood pneumonia cases occur in developing countries, specifically in Africa and Southeast Asia. Pneumonia caused by bacterial pathogens needs urgent medical appointment with immediate antibiotic treatment and caused by viral pathogens needs supportive cares [2]. On the other hand, on the 11th of March, 2020 WHO declares COVID-19 (Coronavirus disease) a pandemic. Alarming rate of spreading, highly infectious attitude and severity of this disease make it a worldwide concern. The main popular screening method for detecting it is polymerase chain reaction (PCR) [3]. The general objective of PCR testing is to detect SARS-CoV-2 RNA from respiratory specimens collected through a variety of resources like nasopharyngeal or oropharyngeal swabs. PCR testing is highly sensitive and gold standard. PCR testing is very time-consuming, laborious, complicated and needs manual processing. An alternative way of detecting COVID-19 is analysing chest radiography images such as X-ray or computed tomography (CT) images. Patients of COVID-19 shows significance abnormality in chest radiology imaging [4, 5]. In epidemic areas, radiology imaging as screening tool might help [6]. Coronavirus disease (COVID‐19) is an infectious disease and looking at the extent of its spread throughout the world, it has been declared as a pandemic by the World Health Organization (WHO) on the 11th of March, 2020 [7].

An experienced radiologist might easily identify normal, viral pneumonia, bacterial pneumonia and newly emerged COVID-19 X-ray images. But requirement of deep knowledge about anatomical principles, pathology and physiology, specially the variety of pathogens make it a complex job for non-experienced personals. Computer-aided diagnostic (CAD) can possibly aid radiologists and physicians to more swiftly and precisely classify radiography images to detect pneumonia and COVID-19 cases. In this paper a deep neural network is trained using transfer learning process to classify pneumonia and COVID-19 from radiology images. Even though deep neural networks fetch efficient models and achieve satisfying results, they are naturally too large, time-consuming and need much complex computational power to be deployed on edge devices like smartphones or embedded sensor nodes. There have been many efforts to compress these networks or so-called models, and a popular and appealing method is knowledge distillation (KD), where a large pre-trained network (teacher model) is used to train a smaller network (student model). In this test the learned soft labels extracted from a pre-trained network over an image input are supplied into the fully connected dense layers to determine the probability of the image over the possible classes. The soft labels are extracted using knowledge distillation process for each image.

Background

Algorithmic Approach to Image Analysis (Traditional and Neural Network)

Traditional algorithmic approach for classifying images involves handcrafted image segmentation and creating shallow neural machine learning classifier to classify according to classes [8]. Instead of this complex and expensive process [9–11], development of deep convolutional neural network layers and models makes the process automatic and more accurate [12, 13]. Every layer of deep convolutional neural network (CNN) consists of multiple convolutions or systematically image analysis filters that eventually produce a feature map for input of the next layer. The network architecture allows to process the input image as a form of pixels and executes the classification accordingly. In modern process of image classification, one classifier could outperform the traditional multiple classifier, and need of handcrafting is also decreased in a significant level. Convolutional neural networks or CNNs are that kind of artificial neural networks that use multiple perceptron or statistical classifiers to analyse inputs (Image or arrays) and learn weights and bases to multiple parts of images and able to separate each other. One benefit of using CNN is it leverages the use of local spatial coherence in the input images, which allow them to have fewer weights as some parameters are shared. So, using CNN is clearly efficient in terms of memory and complexity.

Medical Image Analysis using Neural Networks

Deep learning or convolutional neural networks show a significant increase in performance in the medical image analysis domain for automated object or disease identification and segmentation operation. Recent distinguished work involves but does not limit to an outline analysis on the forthcoming promise of deep learning [14] and an assembly of significant medical applications on.

• cell image segmentation and tracing [15],

• cerebral microbleed detection [16],

• pulmonary nodule detection in CT images [17],

• lymph node and interstitial lung disease detection and classification [18, 19],

• predicting spinal radiological scores [20] extensions of multi-modal imaging segmentation [21, 22],

• automated pancreas segmentation [23],

• and COVID-19 classification using radiology imaging [24, 25].

But on the other hand, it is unclear how well or efficiently the current deep learning techniques will handle multiple thousands of patient data.

Distilling Neural Networks

Neural networks use distributed representation of feature map of input to interpret and processing in hidden layers [26]. These representations are hard to understand and represent in general form. Knowledge distillation (KD) was first applied on deep neural networks by Hinton et al. [27]. KD was first proposed by Bucila et al. [28].

Knowledge distillation is the process to transfer the representations and abilities learned by a large network (teacher network) to any smaller network (student network). Main drawback of knowledge distillation is that it can only be applied for classification tasks, not for regression tasks [29]. KD uses the “dark knowledge” (softened logit output of the bottom output layer of teacher network) that is transferred to student network. This dark knowledge is more than inter-label correlations and one-hot encoding of labels. In case of regression, the deep network interprets the continuous values, which has a tendency of unknown error distribution; as a result, there is no dark knowledge for deep networks trained for regression. But later [30] present their knowledge distillation process for pose regression. Multiple papers try to propose a framework to train the teacher and student network in parallel. Yim et al. [31] and You et al. [32] proposed a shared layer-representation for it. Czarnecki et al. [33] narrowed down the variance between teacher and student derivatives of the loss shared with the discrepancy from teacher predictions. Tarvainen and Valpola [34] proposed an averaging procedure for model weights for the same purpose. Urban et al. [35] trained a network of convolutional neural networks and teach shallow multilayer perceptron as student networks. Sau and Balasubramanian [36] injected noise into teacher logits for making the student network more robust. Employing several teacher networks is always a way to decrease the accuracy difference between teacher and student network. Zhang et al. [37] proposed deep mutual learning so that teacher and student networks can learn side by side during training.

Hinton et al. [27] claim that the success of knowledge distillation is credited to the logit distribution of the incorrect outputs. Furlanello et al. [38] investigated the success of knowledge distillation via gradients of the loss. Lopez-Paz et al. [39] considered the cause behind the efficiency of knowledge distillation from the perception of learning theory by Vapnik [40] by reviewing the estimation error in empirical risk minimization framework.

Other Approaches for Compressing

Obtaining smaller network with least comptonization with respect to accuracy is one of the hottest topics in research. A variety of approaches, other than distillation include shrinking, factorizing or compressing pre-trained deep neural networks. Neural network compression using product quantization [41] network pruning [42–45], connectivity learning [46, 47], hashing [48] and vector quantization-Huffman encoding [49] have been proposed. Also a few factorizations are also proposed to speed up the network’s prediction [50, 51]. Another proved approach is using low bit networks to train the deep learning model [52–54].

Dataset

In this study, mainly three types of X-ray images are used, namely normal chest X-rays, pneumonia X-ray images and COVID-19 chest X-ray images. Normal and pneumonia images are part of dataset and are 1583 and 4273 in number, respectively. That dataset [55] is already tested by Kermany et al. [56] which achieved an accuracy of 92.8%, with a sensitivity of 93.2% and a specificity of 90.1% over pneumonia classification. COVID-19 X-ray images was taken from [57] and are 262 in number. Total number of images is 5964 dividing into three classes. Figure 1 shows examples of each of the three types of images.

Fig. 1.

Sample data cases (normal, pneumonia and COVID-19 cases)

Process

Transfer Learning

In this paper, transfer learning is employed to tackle the problem of detecting the class of chest X-ray images. Due to the immensity and complication of convolutional neural network architectures, scheming and testing models could be expensive and time-consuming. When approaching an image classification, quick and more efficient results can be achieved by applying a technique known as transfer learning. In transfer learning, the weights and convolutional filters that are capable at one task (trained for any classification task), is reused for another task needing only a small amount of retraining and can be trained or tested on a little number of images or data. This includes using a pre-trained neural network architecture with pre-loaded weights, adjusting it to some extent, and then retraining part or the whole model to perform the new task. The filters learned by one task are used to extract features from images that can then be interpreted by the retrained portion of the neural network in order to perform its new task. In this paper, the deep convolutional neural network known as VGG [58] is used to study transfer learning using pre-training over “Imagenet” dataset [59], and weights are used the same as [58].

Teacher Network Architecture

VGG [23] is a deep neural network architecture developed by Karen Simonyan and Andrew Zisserman of the Visual Geometry Group at Oxford. The “16” and “19” in the names (VGG16 and VGG19) refer to the number of weight layers in the network. VGG16 is very successful in the ImageNet competition in 2014. VGG16 architecture has three convolutional layers in the next three stacks, and VGG19 has four convolutional layers in the next three stacks of layers. These are our main feature extractor convolutional architecture of our study. Then the networks are followed by flatten layer. After that it passes through a fully connected dense layer of hidden layers. And last after using dropout of 20%, it passes through a softmax layer to fetch the classification probabilities.

Convolution is a specialized kind of linear operation. In convolutional layer, a matrix named kernel is passed over the input matrix to create a feature map for the next layer. At every location, an element wise matrix multiplication is performed and sums the result onto the feature map. The rectified linear unit activation function (ReLU) is a piecewise linear function that will output the input if it is positive, otherwise it will output zero. Down sampling of convolutional layers is achieved by applying a pooling layer after non-linearity layer. Pooling helps make the representation become approximately invariant to small translations of the input. Softmax or normalized exponential function normalizes the input vector according to probability distribution to the exponentials of inputs.

The input images resolution is (150 × 150 × 3). Every convolutional layer has filters of (3 × 3) shape with ReLU activation. Max-pooling layers would take (2 × 2) matrix. These layers would get the precise locations of the feature map on the activated images. Flatten layer convert every array into one dimensional array. After flattening, the nodes that are achieved from flattening, will act as input layer to fully connected dense layer. These layers also have the same type of rectified activation as convolutional layers. It follows with dropout layer and softmax layer. The model is compiled with Kingma and Ba [60] optimizer, with learning rate 0.0001. Tables 1 and 2 contain the detailed architectures of the both teacher models (VGG16 and VGG19) with their training details and outputs.

Table 1.

VGG16 model architecture

Layer type	Output tensor shape (per image)	Status
Input layer	150, 150, 3	New
Convolutional (2D)	150, 150, 64	Frozen
Convolutional (2D)	150, 150, 64	Frozen
Max-pooling (2D)	75, 75, 64	Frozen
Convolutional (2D)	75, 75, 128	Frozen
Convolutional (2D)	75, 75, 128	Frozen
Max-pooling (2D)	37, 37, 128	Frozen
Convolutional (2D)	37, 37, 256	Frozen
Convolutional (2D)	37, 37, 256	Frozen
Convolutional (2D)	37, 37, 256	Frozen
Max-pooling (2D)	18, 18, 256	Frozen
Convolutional (2D)	18, 18, 512	Frozen
Convolutional (2D)	18, 18, 512	Frozen
Convolutional (2D)	18, 18, 512	Frozen
Max-pooling (2D)	9, 9, 512	Frozen
Convolutional (2D)	9, 9, 512	Frozen
Convolutional (2D)	9, 9, 512	Frozen
Convolutional (2D)	9, 9, 512	Frozen
Max-pooling (2D)	4, 4, 512	Frozen
Flatten	8192	New
Dense/fully connected	4096	New
Dropout (20%)	4096	New
Softmax	3	New

Table 2.

VGG19 model architecture

Layer type	Output tensor shape (per image)	Status
Input layer	150, 150, 3	New
Convolutional (2D)	150, 150, 64	Frozen
Convolutional (2D)	150, 150, 64	Frozen
Max-pooling (2D)	75, 75, 64	Frozen
Convolutional (2D)	75, 75, 128	Frozen
Convolutional (2D)	75, 75, 128	Frozen
Max-pooling (2D)	37, 37, 128	Frozen
Convolutional (2D)	37, 37, 256	Frozen
Convolutional (2D)	37, 37, 256	Frozen
Convolutional (2D)	37, 37, 256	Frozen
Convolutional (2D)	37, 37, 256	Frozen
Max-pooling (2D)	18, 18, 256	Frozen
Convolutional (2D)	18, 18, 512	Frozen
Convolutional (2D)	18, 18, 512	Frozen
Convolutional (2D)	18, 18, 512	Frozen
Convolutional (2D)	18, 18, 512	Frozen
Max-pooling (2D)	9, 9, 512	Frozen
Convolutional (2D)	9, 9, 512	Frozen
Convolutional (2D)	9, 9, 512	Frozen
Convolutional (2D)	9, 9, 512	Frozen
Convolutional (2D)	9, 9, 512	Frozen
Max-pooling (2D)	4, 4, 512	Frozen
Flatten	8192	New
Dense/fully connected	4096	New
Dropout (20%)	4096	New
Softmax	3	New

Student Network Architecture

Student networks are dense neural network models with input dimension same as the logit output shape of the specific layer. The main model is a fully connected dense layer with variable number of filters. In this study, number filters are altered, and the model accuracy is checked. The number of filters in dense network is called the “hidden neuron” count of the student model. Student model size depends upon the filter number (hidden neuron number) and input shape. The accuracy of the model depends upon the accuracy and learning of the teacher model. The dense network has passed through dropout rate of 0.2 to nullify overfitting. At last, the model is compiled using Adadelta [61] optimizer with unity learning rate and 0.95 decay factor.

Distillation

Neural networks typically yield class-probabilities using a softmax layer. Softmax output layer converts the logit, z_i, calculated for each class into a probability, q_i, by comparing z_i with the other logits.

$$q_i={(e}^{\frac{z_i}{T}})/{\sum }_j(e^{\frac{z_j}{T}})$$ 1

where T is a constant known as temperature that is conventionally agreed as the value 1.

In the conventional procedure of distillation, knowledge or learning is shifted to the smaller model by training it on a transfer set. The inputs of the distilled model or smaller model are soft target distribution for each case in the transfer set. Hinton et al. [27] uses a higher temperature value (higher than 1) in softmax to produce soft target distribution of the transfer set that is produced by using the cumbersome model. The same high temperature is retained when training the smaller model, but after it would be trained using a temperature of 1.

Each case in the transfer set contributes a cross-entropy gradient, dC/dz_i, with respect to every logit, z_i of the distilled model. If the cumbersome model has logits v_i which produce soft target probabilities p_i, and the transfer training is done at a temperature of T, this gradient is given by:

$$\frac{\mathit{dC}}{d{z}_i}={(q}_i-p_i)(1/T)=(1/T)({(e}^{\frac{z_i}{T}})/{\sum }_j(e^{\frac{z_i}{T}})-{(e}^{\frac{v_i}{T}})/{\sum }_j(e^{\frac{v_j}{T}}))$$ 2

Using a higher value for T produces a softer probability distribution over classes. For high temperature value, knowledge distillation is equivalent to minimizing 1/2(z_i − v_i)², provided the logits are zero mean separately for each transfer case. At lower temperatures, distillation wages much less consideration to alike logits that are much more negative than the average.

Results

Parameter count of convolutional neural network model signifies the size of the model. Accuracy of the model equals to the ratio of the correctly predicted images and total images in every test. Recall of cases signifies the ratio of correctly predicted images and the total images under that specific case. Covariance is calculated as the average of the product between the values from each sample, where the values had their mean subtracted. The Pearson’s correlation coefficient is determined as the ratio of covariance of the two variables divided and the product of the standard deviation of each data sample. It is the normalization of the covariance between the two variables to give an interpretable score. The total sample of 5964 images is divided into two sets randomly. Four thousand eight hundred 90 images containing all the classes are used for training, and 1074 images are used to test the accuracy of the model in every test case. The main architectures are trained 30 epochs for every different hidden neuron count in the fully connected layer. The main teacher models are trained 30 epochs according to architectures (hidden neuron count varies in fully connected layer) in Tables 1 and 2. The parameter count, accuracy and recall of largest teacher CNNs are mentioned in Table 3. The student architectures are trained 40 epochs for each test. Tables 4 and 5 note all the accuracies of the main architectures (VGG16 and VGG19 respectively) and student distilled architectures with respect to the changing number of hidden layers in the dense layer. Total time taken for the tests on main VGG19 models is 52 min,, 29 s and using distillation process is 8 min, 28 s. Total time taken for the tests on main VGG16 models is 45 min, 28 s and using distillation process with student networks is 8 min, 51 s. Table 6 shows the covariance and Pearson’s correlation in between the accuracies of the two main networks with student accuracies. Figures 2 and 3 show the graphical representation of the main model and student model accuracies for VGG16 and VGG19, respectively. Figures 4 and 5 show these results of the main model accuracy with student model accuracy in scatter plot.

Table 3.

Details of teacher model test

CNN name	CNN characteristics	Parameters (millions)	Accuracy (%)	Recall of normal cases	Recall of pneumonia cases	Recall of COVID-19 cases
VGG16	Dense layer with 4096 hidden neurons	43.29	96.37	0.93	0.97	0.91
VGG19	Dense layer with 4096 hidden neurons	53.6	96.28	0.94	0.96	0.96

Table 4.

Accuracy of different VGG16 models

Number of neurons in fully connected dense layer	Parameters for VGG16-type model (millions)	Accuracy for VGG16-type model (%)	Student accuracy for VGG16-type model (%)
4096	48.29	96.37	96.18
3072	39.89	96.46	96.28
2048	31.50	96.46	96.28
1024	23.10	96.28	95.99
512	18.91	96.46	96.28
256	16.81	96.46	96.28
128	15.76	96.55	96.37
64	15.24	96.65	96.37
32	14.98	96.28	95.99
16	14.85	96.28	95.16
8	14.78	96.18	94.41

Table 5.

Accuracy of different VGG19 models

Number of neurons in fully connected dense layer	Parameters for VGG19-type model (millions)	Accuracy for VGG19-type model (%)	Student accuracy for VGG16-type model (%)
4096	53.60	96.28	95.90
3072	45.20	96.28	96.18
2048	36.81	96.46	96.37
1024	28.42	96.65	96.37
512	24.22	96.46	96.37
256	22.12	96.65	96.37
128	21.07	96.46	96.37
64	20.55	96.65	96.55
32	20.29	96.46	96.28
16	20.16	96.28	94.69
8	20.09	93.76	90.50

Table 6.

Covariance and Pearson’s correlation between accuracies

	Covariance	Pearson’s correlation
VGG16 (main model vs. student model)	0.066	0.774
VGG19 (main model vs. student model)	1.433	0.975

Fig. 2.

Main models and student model accuracy graph for VGG16

Fig. 3.

Main model and student model accuracy graph for VGG19

Fig. 4.

Main model vs. student model accuracy (VGG16) scatter plot

Fig. 5.

Main model vs. student model accuracy (VGG19) scatter plot

Discussion and Conclusion

Table 6 covariance is positive in every case, so the accuracies of main model and student model are changing in the same direction. Pearson’s correlation value for both is more than 0.5. This suggests a high level of correlation, e.g. a value above 0.5 and close to 1.0. So, the accuracy of the student and main architecture is slightly different from each other, but linearly correlate-able with each other. That is how we can find the best hidden neuron number for our dense layer using this technique that is saving nearly 44 min for VGG19 and 36 min, 37 s for VGG16 feature extractor. We can use any pre-trained feature extractor (even if it is new architecture built for any specific task or trained upon any other dataset than “ImageNet”) and find the optimal number of hidden layer in top dense layer as it performs the best on that task. Determining the best number of neurons, not only shorten the memory necessity of the deep learning model but also increases the accuracy, by nullifying the option of overfitting over that data. Using the same task manually, as the ideal number of hidden neurons will differ from architecture to architecture and dataset to dataset, will take a long time. But this process does not give the perfect best accuracy; the accuracy of student models is slightly less than that of the main model; but as the study suggests, these student accuracies will help to roughly interpret the best accuracy. On the other hand, this process might be slow in nullifying the overfitting than pruning, but it offers better flexibility and better relative accuracy than pruning as it helps to rebuild the whole network the best way possible. But this process can delete the overfitting on the top layers; pruning can delete overfit for multiple layers at once.

Declarations

Conflict of Interest

The author declares no competing interests.