Harmonization and Targeted Feature Dropout for Generalized Segmentation: Application to Multi-site Traumatic Brain Injury Images
While learning based methods have brought extremely promising results in medical imaging, a major bottleneck is the lack of generalizability. Medical images are often collected from multiple sites and/or protocols for increasing statistical power, while CNN trained on one site typically cannot be well-transferred to others. Further, expert-defined manual labels for medical images are typically rare, making training a dedicated CNN for each site unpractical, so it is important to make best use of the limited labeled source data. To address this problem, we harmonize the target data using adversarial learning, and propose targeted feature dropout (TFD) to enhance the robustness of the model to variations in target images. Specifically, TFD is guided by attention to stochastically remove some of the most discriminative features. Essentially, this technique combines the benefits of attention mechanism and dropout, while it does not increase parameters and computational costs, making it well-suited for small neuroimaging datasets. We evaluated our method on a challenging Traumatic Brain Injury (TBI) dataset collected from 13 sites, using labeled source data of only 14 healthy subjects. Experimental results confirmed the feasibility of using the Cycle-consistent adversarial network for harmonizing multi-site MR images, and demonstrated that TFD further improved the generalization of the vanilla segmentation model on TBI data, reaching comparable accuracy with that of the supervised learning. The code is available at https://github.com/YilinLiu97/Targeted-Feature-Dropout.git.
This work was supported by NARSAD: Brain and Behavior grant 24103 (to BN) and National Institutes of Health grant funding NINDS R01 NS092870, NIMH P50 MH100031 and a core grant to the Waisman Center from the National Institute of Child Health and Human Development (U54 HD090256). Disclosure Statement: A Alexander is part owner of Thervoyant, Inc. We gratefully acknowledge the support of NVIDIA Corporation with the donation of the Titan Xp and the Telsa K80 used for this research.
- 1.Ghiasi, G., Lin, T.Y., Le, Q.V.: DropBlock: a regularization method for convolutional networks. In: Advances in Neural Information Processing Systems, pp. 10727–10737 (2018)Google Scholar
- 2.Gibson, E., et al.: Inter-site variability in prostate segmentation accuracy using deep learning. In: Frangi, A.F., Schnabel, J.A., Davatzikos, C., Alberola-López, C., Fichtinger, G. (eds.) MICCAI 2018. LNCS, vol. 11073, pp. 506–514. Springer, Cham (2018). https://doi.org/10.1007/978-3-030-00937-3_58CrossRefGoogle Scholar
- 3.He, Y., Kang, G., Dong, X., Fu, Y., Yang, Y.: Soft filter pruning for accelerating deep convolutional neural networks. arXiv preprint arXiv:1808.06866 (2018)
- 4.Hu, J., Shen, L., Sun, G.: Squeeze-and-excitation networks. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 7132–7141 (2018)Google Scholar
- 6.Karani, N., Chaitanya, K., Baumgartner, C., Konukoglu, E.: A lifelong learning approach to brain MR segmentation across scanners and protocols. In: Frangi, A.F., Schnabel, J.A., Davatzikos, C., Alberola-López, C., Fichtinger, G. (eds.) MICCAI 2018. LNCS, vol. 11070, pp. 476–484. Springer, Cham (2018). https://doi.org/10.1007/978-3-030-00928-1_54CrossRefGoogle Scholar
- 7.Lin, M., Chen, Q., Yan, S.: Network in network. arXiv preprint arXiv:1312.4400 (2013)
- 8.Roy, A.G., Navab, N., Wachinger, C.: Concurrent spatial and channel squeeze & excitationin fully convolutional networks. In: Frangi, A., Schnabel, J., Davatzikos, C., Alberola-López, C., Fichtinger, G. (eds.) MICCAI 2018. LNCS, pp. 421–429. Springer, Cham (2018). https://doi.org/10.1007/978-3-030-00928-1_48CrossRefGoogle Scholar
- 9.Shu, R., Bui, H.H., Narui, H., Ermon, S.: A DIRT-T approach to unsupervised domain adaptation. arXiv preprint arXiv:1802.08735 (2018)
- 10.Tompson, J., Goroshin, R., Jain, A., LeCun, Y., Bregler, C.: Efficient object localization using convolutional networks. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 648–656 (2015)Google Scholar
- 11.Zhou, B., Khosla, A., Lapedriza, A., Oliva, A., Torralba, A.: Learning deep features for discriminative localization. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 2921–2929 (2016)Google Scholar
- 12.Zhu, J.Y., Park, T., Isola, P., Efros, A.A.: Unpaired image-to-image translation using cycle-consistent adversarial networks. In: Proceedings of the IEEE International Conference on Computer Vision, pp. 2223–2232 (2017)Google Scholar