This blog post shows the findings of a brain tumor classification problem, which was conducted for the talk "Machine Learning in medicine" for the ENEA 2020 conference (the yearly conference of the European Neuroendocrine Association).
Hardware
For this problem a NVidia Jetson TX2 developer kit was used. Probably that board was more designed to be used for inference on the edge, but it turned out to be fast enough to train the whole model.
The reason for doing this with the board instead of a desktop computer with even better GPU was that when trying to use GPU-accelerated machine learning it can be tricky to get all the software package installation right: E.g. the combination of Ubuntu / NVidia's cuDNN library / GPU-accelerated Tensorflow / Keras needs specific versions of each package and might or might not work with other versions.
The Jetson developer board comes pre-installed with working software, and other needed parts like pythons' tensorflow-gpu package can be installed from a desktop system with NVidia's Jetpack installer (as of this writing that installer needs Ubuntu 18.04 and will not run on 20.04, so it had to be used from inside a VMPlayer instance).
Data set
Kaggle is a great resource for free data sets with interesting problems to learn from. The data set used here is "Brain Tumor Classification (MRI)" and consists of ~ 3200 MRI images, each one belonging to four different tumor classes (see above): Glioma, meningioma, pituitary tumor and no tumor.
This data set comes with the data already split into a training and a test set, each of them containing one subdirectory for each tumor category. So when using the right Keras functions, there is little preparation to do with the data before training the model.
Machine learning model
The 3200 images from the data set is already quite a high number, but not something that would guarantee a good accuracy without overfitting. However, data augmentation, i.e. artificially creating more data by flipping, rotating etc. the training data, did not help that much in this case, interestingly.
A rather recent trend in machine learning is transfer learning: A model that was used to train other images is used for the problem at hand, and while training for the new problem (in our case: MRI images), the weights of the old model can be used and carefully adapted.
The learning rate here should be low, and depending on the similarity of the models, a specific amount of layers should be frozen. These articles provide some further reading: Transfer learning and fine-tuning and A Simple Guide to Using Keras Pretrained Models.
The code for the brain tumor classification is using Keras and Tensorflow with weights pre-trained from another model (ImageNet).
The code can be found here: https://github.com/edelhirsch/brain-tumor-classification
Results
When using transfer learning, the model doesn't need many epochs to achieve a good accuracy and low loss: In the case at hand, something around 10 epochs or even less is enough, see below. There is a better accuracy of 97.6% after 14 epochs, but this also seems to pose a bigger risk on overfitting the model.
Evaluating the model with the test data shows that there is still room for improvement: The accuracy is 79.4% with a loss of 1.3. This hints at some overfitting during the training.
Next steps
There are still some things left to do:
It would be nice to turn the code above into a jupyter notebook and upload it to the Kaggle site, happening hopefully soon (™).
It seems unlikely that another image model might provide better results, but might be worth a try. There are several ones available on Tensorflow hub.
Since the images from the original model and the MRI images are quite different, it would be interesting to try generating a model from other MRI images and transfer the weights from that model.
Hints on how to further improve the code or merge requests are always appreciated.
Do you have questions or a similar problem to solve? Contact us or book a free meeting!