Transfer learning for efficient classification in lung cancer detection using mobilenet architecture

Author: Trisha Sengupta

Introduction and Motivation

Carcinoma is a type of cancer that grows in skin or on the tissues of organs, also known as the epithelial tissue. There are two types of “non small cell” cancers associated with carcinoma. These malignant cancers are called squamous cell carcinoma and adenocarcinoma. Detection of this cancer, as well as the type of the cancer is an important and challenging problem. Correct classification of the cancer type is important because the targeted treatment is different for each type.

One possible way to detect if the cancer is present is to refer to a pathologist, who usually studies body fluid and tissue to identify diseases. Pathologists however cannot tell apart the type of non small cell cancer, as this is difficult to do with a human eye. This is the primary reason why we need to adopt a Machine Learning based system. We adopt the classification strategy using deep convolutional neural networks, similar to strategy adopted in a prior work in this area [1]. However, our focus is to make all this work on low end devices, and make our training faster and usable with a small amount of data.

Impact

Our research has 2 major impacts:

1. The work focuses on using a mobile friendly neural network architecture that can work on the “edge” devices like mobile phone, or even a chrome browser. This helps us to diagnose CT scans in the field where there may not be heavy duty computers with GPU or TPU available.
2. The work also shows how we can use the concept of transfer learning: where we use a model that was pre-trained for a general image classification tasks, but only a “small” part of it was tuned for our specific task. This is impactful because it reduces the burden on the data to be collected. Data collection is always an expensive and tedious process.

Background & Challenges

As mentioned earlier, in this research, we use deep learning to classify a cancer cell into squamous cell carcinoma and adenocarcinoma. Unlike the technique described in [1], our aim is to use a very lean architecture of the neural network; so that this can run on very inexpensive hardware on the “edge”, such as mobile devices and laptops.

Requirements

• Data: https://github.com/tampapath/lung_colon_image_set. The dataset is described in this paper [2, 4].
• Dependent software/tools: Google Colab (install), Tensorflow, Octave/Matlab.
• Hardware: Computer with GPU.

Data and augmentation method

Our datasets contained 5000 images of lung squamous cell carcinoma and 5000 images of lung adenocarcinoma from the LC25000 dataset. We added augmentations such as random noise and flips to add more variety. We still note that the number of images is not nearly enough to train a large network with millions of parameters. Such large networks usually overfit the problem we are trying to solve. Hence we explored concepts of transfer learning, freezing most of the parameters while exposing only a few handful parameters as well as starting with a very lean network, as explained below.

Figure 1: Examples of the training images.

ML Models

The ML model we adopt here is a Mobilenet V2 [5] with depth multiplier 0.5. We only make around 164K of a total of 870K parameters trainable. Compared to the method in [1, 3], which had 11M parameters to train, our model size is 12 times smaller. Also, our number of trainable parameters are approximately 67 times less.

(a)

(b)

Figure 2: (a) The mobilenet backbone used in our research vs. (b) the larger backbone used in [1].

Training and Evaluation Experiments

Our training is very swift compared to the one in [1], mainly due to the fact that the number of training parameters are 67 times less. The training and validation accuracy as well as the overall loss can be seen in Figure 3.

Figure 3: Training/validation accuracy and loss using our model.

Our test classification accuracy is 90.8%, just 3.2% less than the accuracy of 94.0% in [1] while the number of training parameters is 67 times less.

Our code can be found in github at this location: https://github.com/trishas2022/lungs

Conclusions

Our experimental results show that for lung cancer classification tasks using a mobile friendly model and transfer learning, we can use a small amount of training corpus and still achieve accuracy closer to that of a large model. Our research findings point to the fact that the method is deployable on computers and mobile platforms that are much leaner compared to a server with dedicated GPUs.

Reference

1. Sub-classifying Lung Cancer with TensorFlow 2 and Keras.
2. New Freely Available Lung and Colon Cancer Image Dataset for ML Researchers
3. https://arxiv.org/abs/1912.12142
4. Dataset: https://github.com/tampapath/lung_colon_image_set
5. MobileNetV2: Inverted Residuals and Linear Bottlenecks: https://arxiv.org/abs/1801.04381
6. Our code: https://github.com/trishas2022/lungs