Combined Transfer Learning and Test-Time Augmentation Improves Convolutional Neural Network-Based Semantic Segmentation of Prostate Cancer from Multi-Parametric MR Images

Abstract

PurposeMultiparametric MRI (mp-MRI) is a widely used tool for diagnosing and staging prostate cancer. The purpose of this study was to evaluate whether transfer learning, unsupervised pre-training and test-time augmentation significantly improved the performance of a convolutional neural network (CNN) for pixel-by-pixel prediction of cancer vs. non-cancer using mp-MRI datasets. Methods154 subjects undergoing mp-MRI were prospectively recruited, 16 of whom subsequently underwent radical prostatectomy. Logistic regression, random forest and CNN models were trained on mp-MRI data using histopathology as the gold standard. Transfer learning, unsupervised pre-training and test-time augmentation were used to boost CNN performance. Models were evaluated using Dice score and area under the receiver operating curve (AUROC) with leave-one-subject-out cross validation. Permutation feature importance testing was performed to evaluate the relative value of each MR contrast to CNN model performance. Statistical significance (p<0.05) was determined using the paired Wilcoxon signed rank test with Benjamini-Hochberg correction for multiple comparisons. ResultsBaseline CNN outperformed logistic regression and random forest models. Transfer learning and unsupervised pre-training did not significantly improve CNN performance over baseline; however, test-time augmentation resulted in significantly higher Dice scores over both baseline CNN and CNN plus either of transfer learning or unsupervised pre-training. The best performing model was CNN with transfer learning and test-time augmentation (Dice score of 0.59 and AUROC of 0.93). The most important contrast was apparent diffusion coefficient (ADC), followed by Ktrans and T2, although each contributed significantly to classifier performance. ConclusionsThe addition of transfer learning and test-time augmentation resulted in significant improvement in CNN segmentation performance in a small set of prostate cancer mp-MRI data. Results suggest that these techniques may be more broadly useful for the optimization of deep learning algorithms applied to the problem of semantic segmentation in biomedical image datasets. However, further work is needed to improve the generalizability of the specific model presented herein.