Evaluating TPUs and GPUs in a Two-View EfficientNet-based architecture for cancer classification on mammograms: performance and speed analysis

Introduction

Breast cancer is the most prevalent cancer among women worldwide. Mammography is the primary exam used to detect this disease in its early stages. Currently, radiologists interpret these radiological images, but CAD (Computer-Aided Detection and Diagnosis) systems have been developed to assist in this process. While GPUs have traditionally been used for training these systems, newer hardware like TPUs (Tensor Processing Units) has been designed specifically for machine learning tasks, and offers advantages over GPUs that can be explored, such as having more memory.

Methods

This work compared the performance of a two-view mammogram classifier proposed by Daniel Petrini et al. in "Breast Cancer Diagnosis in Two-View Mammography Using End-to-End Trained EfficientNet-Based Convolutional Network" and its components (the one-view classifier and patch classifier) on the public dataset CBIS-DDSM (Curated Breast Imaging Subset of the Digital Database for Screening Mammography). The comparison was made using both GPUs and TPUs, leveraging the extra memory and specialized architecture of TPUs.

Results

Training on TPUs was up to 18 times faster than on GPUs, a significant increase in training speed, potentially leading to better models in future work. However, no conclusive evidence showed that using higher resolution images with TPUs improved model performance. Metrics (accuracy and ROC-AUC) were similar at 1152x896 (GPU) and at 2304x1792 (TPU).

Conclusions

Although the classification performance did not improve when increasing exam resolution, the use of TPUs is justifiable due to the increase in training speed, opening up possibilities to train with more data and using more complex architectures, which could lead to better classification results.