High-Accuracy Covid 19 Prediction from Chest X-Ray Images using Pre-Trained Convolutional Neural Networks in PyTorch | by Bharat Sethuraman Sharman | Towards Data Science

Leveraging Efficientnet architecture to achieve 99%+ prediction accuracy on a Medical Imaging Dataset pertaining to Covid19
In this post, I will share my experience of developing a Convolutional Neural Network algorithm to predict Covid-19 from chest X-Ray images with high accuracy. I developed this algorithm while participating in an In-Class Kaggle competition for a Ph.D. level course on Deep Learning. I was happy to learn from the final competition Leaderboard that I stood 1st in the competition amongst 40 participating Ph.D. and Masters students of Statistics and Computer Science departments! Instead of moving straight to results, I will mention the challenges that I faced while I developed and trained the model and how I found ways to resolve them. While working on this problem, I experienced the truth first-hand in the saying “Enjoy the process, the result will take care of itself”!
You can find my complete code on GitHub here.
The Challenge in Brief
Given an input chest X-ray image, the algorithm must detect whether the person has been infected with Covid-19 or not. We have a set of supervised X-ray images that have been carefully labeled by radiologists on which the model must be trained. Such a model can help health care professionals diagnose COVID-19 cases much more quickly than a radiologist having to go through each scan one by one especially in situations where a large number of people need to be tested in a short time window.
Let’s take an example. The images shown below from the competition website are of people without (negative) and with (positive) Covid-19.
Dataset Description
The training dataset consists of 15264 (512x512 pixel) images that have been classified by an expert radiologist. The test dataset consists of 400 such images coming from the same distribution. You can find the dataset here.
I found that the training dataset is significantly imbalanced with about 90% of the images being of the non-Covid class. The problem with significantly imbalanced sets is that even a naive learning algorithm that simply outputs the class of the majority class as the output will attain high accuracy. In other words, it will label everyone as being Covid-free and attain an ‘accuracy’ of 90% even though about 10% of the people actually have Covid. Initially, I started with the imbalanced training dataset as it is. From the benefit of hindsight, it turned out to be not the best starting point but one must start with what one has and in this case, it was the training dataset. I will describe it more in detail later in this article.
Climbing the Performance Ladder — Steps and Missteps
A deep learning algorithm has many moving parts and I found that training a high-performing model is as much an art as it is math. One does not know most of the answers when starting out and has to learn by experimentation.
I will now describe how I went about selecting each of the hyperparameters and what I learned in the process.
Learning Rate: A lower learning rate leads to better convergence but not necessarily the best test accuracy. The issue could be that the algorithm has not yet learned the set of weights that lead to the lowest error for the test set. A higher learning rate, on the other hand, can lead to the algorithm overshooting the optimal set of weights. I experimented with learning rates of 10e-3, 10e-4, and 10e-5 and found that 10e-4 led to the lowest error rates for the test dataset.
Batch Size: On increasing the batch size to 256 or even down to 64, I ran into memory issues and therefore had to opt for a lower batch size. I experimented with 32 and 16. Out of these two batch sizes, the batch size of 32 led to the lowest error rates with a learning rate of 10e-4. The interplay of Learning Rate, Batch Size, and Accuracy has been explored well in this paper. The authors experimented with various learning rates and batch sizes for the VGG16 CNN. I started with a lower learning rate (lr=0.0001 as aginst 0.001) that turned out to the better choice later and went with the rule of thumb of setting the batch size as 32 which is quite common in deep learning practice. I thank the authors for the leads that I got from their study.
A complex field like Deep Learning is never a lone wolf endeavor and we all stand on the shoulders of each other, not just the giants, to look ahead.
Pre-Trained Model: Coming to the question of which pre-trained model to use; I came across an interesting article on the Google AI Blog that compared the accuracy on the Imagenet dataset vs the number of parameters in the model for several pre-trained models. Computer Vision is a quickly evolving field and new models keep getting developed at a fast pace. Initially, I had plans to use VGG16 or Resnet50 but when I saw this article on Efficientnet, I changed my mind in the favor of Efficientnet. I have an RTX 2070 GPU and that the limitations of my hardware placed a constraint on the model that I could use. I started with the b7 version of Efficientnet but my GPU quickly gave up as the model was not fitting in its memory. After numerous Kill -9 job_number on b7,b6,b5, and b4 models(the command to delete a model from computer memory in Linux :) ), I settled with Efficientnet b3. The sole reason to do so at that point in time was that it was the largest model in the Efficientnet family that finally fit within my GPU memory. Later on, this turned out to be an excellent choice!
Epochs: With a model like Efficientnet b3 that has 12 Million parameters and the number of training dataset instances being in the order of tens of thousands, it is very easy to overtrain a model. In fact, I got a training accuracy of 1 while training the model with 20 Epochs! Don’t ask me the test accuracy for that case :) I used epochs ranging from 1 to 30 for various learning rate and batch size combinations.
Sampling: I started with no sampling. Strangely, the model was classifying almost all the test cases as “1”, contrary to my expectation of “0”. I felt that the problem was that the test images were not being transformed as the training images were and the untransformed test images were somehow resembling the transformed training images of label “1” (more on this crucial topic later). I then used a weighted random sampler that led to an almost 50–50 representation of the 0 and 1 classes in the training dataset. However, it also led to the overall training dataset size reducing from ~15000 to just under ~3000. The accuracy improved but plateaued at around 0.86 with this sampling technique. Next, I oversampled from the minority class and brought the number of images from both the classes on par with each other. By doing so, my training dataset increased from 15000 to 26000. This certainly helped as I was now hitting the 90s in accuracy and I got an accuracy of about 0.94.
Transformations- Appling the same transformations used on the training dataset to the test dataset: I was plateauing at 0.94 and others were getting accuracies of up to 0.975. This set me thinking hard about what I was missing in my model. I traced and re-traced the steps several times, while bathing, cooking, and walking.
Suddenly, the Eureka moment came: Distribution! All machine learning rests on the assumption that the distribution of the test data is the same as that of the training data. Otherwise, what we are learning from the training dataset simply will not apply to the test dataset. By transforming the training images via rotations and translations, I was changing the distribution. However, by leaving the test images as it is, I was inadvertently applying the algorithm on a “different” distribution. I immediately applied identical transformations on the test images and lo and behold, I saw an accuracy of 0.975 for the first time!
Tuning: Once this insight came, all that was left to be done was tuning the hyperparameters to achieve the best possible result. I tried various epochs and found that for this dataset with the pre-trained model, learning rate, and batch size that I used, I got the best test accuracy after 11 epochs of training.
Could have I done better?
Perhaps using Efficientnet-b7 coupled with a larger batch size could have led to even higher accuracy. If someone has a more powerful GPU/distributed computing system, please give it a shot and let me know.
Practical Application: With this algorithm, a medical practitioner can identify Covid accurately in 99 out of 100 people rapidly (a matter of 40 minutes to train the algorithm on 15000 images and run it on 400 images) Taking into account only the classification time for the test images, it takes merely 1 minute for the algorithm to classify 400 images. So this can be very useful for field-staff to rapidly screen images and identify people who test positive.
Having walked you through the algorithm discovery journey, I now move to the nuts-and-bolts aka development.
Algorithm Development in PyTorch
I developed my algorithm in PyTorch which is a Machine Learning Library that is well-suited for Computer Vision tasks.
As the first step, I imported numpy, pandas, os, torch, etc, and set the base path where the data folder is located. Then I created pandas data frames to contain the training dataset image names and their labels and the test dataset image names respectively.
As the first step, I imported numpy, pandas, os, torch, etc, and set the base path where the data folder is located. Then I created pandas data frames to contain the training dataset image names and their labels and the test dataset image names respectively.
2. To check whether the data frames were loaded properly, it is good to check the first few rows of train_dataset and the test_dataset. Here is the code and the resulting output from the Jupyter notebook:
We see that the names of the images are in the File column and the Labels are in the Label column for the training images. For the test images, we have the names in the File column.
3. Here’s the code to view some of the images in the training dataset folder,
This code will result in 15 X-Ray images of the type shown above.
4. The next step is to create the training dataset out of the images and the labels in a form that Pytorch can process. Dataloader is a function in Pytorch that enables us to conveniently pair images and their labels to create a matrix/tensor that can then be used by PyTorch as an input to a learning algorithm.
5. I then set the learning rate to be 10e-4
6. Next, I prepared the training set using the Dataset dataloader class that I created in step 4. I named it as training_set_untransformed as I have not yet applied the data transformations like rotation or translation to the images.
7. I then created a transforms.Compose method to resize the images, randomly apply rotation and horizontal translation and convert the resulting matrix to a tensor. Such a transform is useful for data augmentation.
8. The next step is to use the transforms method on the training dataset. For the images in the minority class (covid cases), I have upsampled them so that the number of images in both the classes are the same after the transformation is applied.
9. I then split these newly created images into training and validation sets in the 80:20 ratio.
10. Using the Dataloader function of Pytorch, I created batches of 32 each from the training dataset and shuffled them to ensure roughly equal representation of both classes in all the batches.
11. To enable GPU computation, I used the cuda.is_available() command. The Deep Learning code can also run on a CPU but would take much longer.
12. The following is the command to import Efficientnet-b3 from Pytorch and set the number of classes to 2.
13. After importing, the model has to be loaded to the memory of the GPU. If your model is too big for the GPU memory, you will get an error message and would need to select another model that will fit into the memory. Run the command on a Jupyter notebook and admire the complexity of Efficientnet!
14. It is good to save the weights generated by the training process periodically. In case if the model crashes due to some reason, we can call the saved weights and resume training instead of having to start from scratch again.
15. I used the cross-entropy loss as the criterion for loss calculation as we are dealing with a classification problem. I used the Adam optimizer, which is a well-known weight update method for deep learning, and set the number of epochs to 11 (this learning came after a lot of experimentation). Learning rate decay is a parameter used to update the learning rate after each epoch and I set a conservative decay rate of 0.99 as the learning rate is already low. I created two lists to track accuracy and loss histories.
16. Training the model: Here’s the training part of the code. I ran the model for 11 epochs and saved the 1st, 10th, and final (11th) versions of the model. It took about 40 minutes to train it.
At the end of the 11th Epoch, I got an accuracy of 99% on the training dataset for both classes. Accuracy of     0: 99 %Accuracy of     1: 99 %[11 epoch] Accuracy of the network on the Training images: 99 
17. Here is a plot of the training accuracy and cross-entropy loss over iterations. There are 7227 batch iterations and the plot shows how these values are changing over the batch iterations.
18. The next step is to apply the transforms. Compose to the test dataset as well to make it conform to the same distribution as the training dataset.
19. To predict the class of test images, we need to define a predict_image function:
20. The following code snippet checks how well the model is doing on the validation dataset.
21. With the following lines, one can predict the classes of the test set images and store them in a csv file.
Conclusion
Through this Kaggle Competition, I learned that conceiving and developing a Deep Learning Algorithm takes some patience and we must continuously keep looking for chinks in the model as well as the discoveries that others have made to make our models better. My article is a small way of paying it forward to the Machine Learning community so that we can share and learn together and improve our algorithms to solve problems in better ways!
Hope you could learn something useful from my post. Please share your valuable feedback!
Acknowledgement: I gratefully acknowledge Kartik Dutt who had developed code for a different Computer Vision challenge and whose parts of whose code I found appropriate for this deep learning task and incorporated in my work. (https://github.com/kartikdutt18/Kaggle-ATPOS-Efficient_Net/blob/master/Efficient_Net.ipynb)