Evolutionary pruning of transfer learned deep convolutional neural network for breast cancer diagnosis in digital breast tomosynthesis

Abstract

Deep learning models are highly parameterized, resulting in difficulty in inference and transfer learning for image recognition tasks. In this work, we propose a layered pathway evolution method to compress a deep convolutional neural network (DCNN) for classification of masses in digital breast tomosynthesis (DBT). The objective is to prune the number of tunable parameters while preserving the classification accuracy. In the first stage transfer learning, 19 632 augmented regions-of-interest (ROIs) from 2454 mass lesions on mammograms were used to train a pre-trained DCNN on ImageNet. In the second stage transfer learning, the DCNN was used as a feature extractor followed by feature selection and random forest classification. The pathway evolution was performed using genetic algorithm in an iterative approach with tournament selection driven by count-preserving crossover and mutation. The second stage was trained with 9120 DBT ROIs from 228 mass lesions using leave-one-case-out cross-validation. The DCNN was reduced by 87% in the number of neurons, 34% in the number of parameters, and 95% in the number of multiply-and-add operations required in the convolutional layers. The test AUC on 89 mass lesions from 94 independent DBT cases before and after pruning were 0.88 and 0.90, respectively, and the difference was not statistically significant (p > 0.05). The proposed DCNN compression approach can reduce the number of required operations by 95% while maintaining the classification performance. The approach can be extended to other deep neural networks and imaging tasks where transfer learning is appropriate.

Fig. 1

Flow chart of the proposed two-stage transfer learning and evolutionary pruning approach. (left) AlexNet trained on ImageNet. (middle) DCNN transfer network trained on mammography. (right) Transfer learning of DCNN for classification in DBT.

Fig. 2

Data set characteristics. The number of views consists of craniocaudal (CC) and mediolateral (MLO) views, including CC and MLO views in the mammography data sets, CC and MLO views in the DBT-UM set and MLO views in the DBT-MGH set. Note that each patient case can have a single or multiple lesions and a lesion may not be visible on all views. After data augmentation, mammography: 19,632 ROIs and DBT: 12,680 ROIs.

Fig. 3

Layered pathway evolution using genetic algorithm. The table at the top shows the number of initial active nodes of the AlexNet (leftmost) from C₁ to C₅ (top to bottom) and the number of active nodes at each step of pruning in the convolutional layers. The corresponding population evolution in each convolutional layer is shown in the graph below. Five parallel bars at the top indicate the five convolutional layers either pruning, frozen or pruned. The ROI-based validation AUCs from the LOOCV resampling method are shown.

Fig. 4

Filter kernels of size 11 × 11 pixels from convolutional layer, C₁. Top: 64 nodes from original DCNN, bottom: 16 nodes from pruned DCNN. C₁ layer is frozen during the first stage of transfer learning.

Fig. 5

ROC curves for independent DBT test between original DCNN and pruned DCNN.