Skip to main content

Memory Efficient 3D U-Net with Reversible Mobile Inverted Bottlenecks for Brain Tumor Segmentation

Part of the Lecture Notes in Computer Science book series (LNIP,volume 12659)

Abstract

We propose combining memory saving techniques with traditional U-Net architectures to increase the complexity of the models on the Brain Tumor Segmentation (BraTS) challenge. The BraTS challenge consists of a 3D segmentation of a 240 \(\times \) 240 \(\times \) 155 \(\times \) 4 input image into a set of tumor classes. Because of the large volume and need for 3D convolutional layers, this task is very memory intensive. To address this, prior approaches use smaller cropped images while constraining the model’s depth and width. Our 3D U-Net uses a reversible version of the mobile inverted bottleneck block defined in MobileNetV2, MnasNet and the more recent EfficientNet architectures to save activation memory during training. Using reversible layers enables the model to recompute input activations given the outputs of that layer, saving memory by eliminating the need to store activations during the forward pass. The inverted residual bottleneck block uses lightweight depthwise separable convolutions to reduce computation by decomposing convolutions into a pointwise convolution and a depthwise convolution. Further, this block inverts traditional bottleneck blocks by placing an intermediate expansion layer between the input and output linear 1 \(\times \) 1 convolution, reducing the total number of channels. Given a fixed memory budget, with these memory saving techniques, we are able to train image volumes up to 3x larger, models with 25% more depth, or models with up to 2x the number of channels than a corresponding non-reversible network.

Keywords

  • Depthwise separable convolution
  • Inverted residual
  • Reversible network

M. Pendse and V. Thangarasa—Equal contribution.

This is a preview of subscription content, access via your institution.

Buying options

Chapter
USD   29.95
Price excludes VAT (USA)
  • DOI: 10.1007/978-3-030-72087-2_34
  • Chapter length: 10 pages
  • Instant PDF download
  • Readable on all devices
  • Own it forever
  • Exclusive offer for individuals only
  • Tax calculation will be finalised during checkout
eBook
USD   89.00
Price excludes VAT (USA)
  • ISBN: 978-3-030-72087-2
  • Instant PDF download
  • Readable on all devices
  • Own it forever
  • Exclusive offer for individuals only
  • Tax calculation will be finalised during checkout
Softcover Book
USD   119.99
Price excludes VAT (USA)
Fig. 1.
Fig. 2.
Fig. 3.
Fig. 4.

References

  1. Bakas, S., et al.: Identifying the best machine learning algorithms for brain tumor segmentation, progression assessment, and overall survival prediction in the brats challenge. ArXiv abs/1811.02629 (2018)

    Google Scholar 

  2. Bakas, S., et al.: Advancing the cancer genome atlas glioma MRI collections with expert segmentation labels and radiomic features. Sci. Data 4, 1–13 (2017)

    CrossRef  Google Scholar 

  3. Bakas, S., et al.: Segmentation labels and radiomic features for the pre-operative scans of the TCGA-GBM collection. Cancer Imaging Arch. (2017). https://doi.org/10.7937/K9/TCIA.2017.KLXWJJ1Q

    CrossRef  Google Scholar 

  4. Bakas, S., et al.: Segmentation labels and radiomic features for the pre-operative scans of the TCGA-LGG collection. Cancer Imaging Arch. (2017). https://doi.org/10.7937/K9/TCIA.2017.GJQ7R0EF

    CrossRef  Google Scholar 

  5. Bauer, S., Wiest, R., Nolte, L.P., Reyes, M.: A survey of MRI-based medical image analysis for brain tumor studies. Phys. Med. Biol. 58, R97–R129 (2013)

    Google Scholar 

  6. Brügger, R., Baumgartner, C.F., Konukoglu, E.: A partially reversible U-net for memory-efficient volumetric image segmentation. In: Shen, D. (ed.) MICCAI 2019. LNCS, vol. 11766, pp. 429–437. Springer, Cham (2019). https://doi.org/10.1007/978-3-030-32248-9_48

    CrossRef  Google Scholar 

  7. Gomez, A.N., Ren, M., Urtasun, R., Grosse, R.B.: The reversible residual network: Backpropagation without storing activations. In: Advances in Neural Information Processing Systems 30, pp. 2214–2224. Curran Associates, Inc. (2017)

    Google Scholar 

  8. Hanif, F., Muzaffar, K., Perveen, k., Malhi, S., Simjee, S.: Glioblastoma multiforme: a review of its epidemiology and pathogenesis through clinical presentation and treatment. Asian Pac. J. Cancer Prev. 18(1), 3–9 (2017)

    Google Scholar 

  9. He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In: 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR, pp. 770–778 (2016)

    Google Scholar 

  10. Iglovikov, V., Shvets, A.: Ternausnet: U-net with vgg11 encoder pre-trained on imagenet for image segmentation (2018)

    Google Scholar 

  11. Isensee, F., et al.: No new-net. In: Crimi, A., van Walsum, T., Bakas, S., Keyvan, F., Reyes, M., Kuijf, H. (eds.) Brainlesion. pp. 234–244. Lecture Notes in Computer Science (including subseries Lecture Notes in Artificial Intelligence and Lecture Notes in Bioinformatics), Springer (January 2019), 4th International MICCAI Brainlesion Workshop, BrainLes 2018 held in conjunction with the Medical Image Computing for Computer Assisted Intervention Conference, MICCAI 2018; Conference date: 16–09-2018 Through 20–09-2018

    Google Scholar 

  12. Kamnitsas, K., et al.: Deepmedic for brain tumor segmentation. In: MICCAI Brain Lesion Workshop (October 2016)

    Google Scholar 

  13. Long, J., Shelhamer, E., Darrell, T.: Fully convolutional networks for semantic segmentation. In: 2015 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 3431–3440 (2015)

    Google Scholar 

  14. Menze, B.H., Jakab, A., Bauer, S., Kalpathy-Cramer, J., Farahani, K., et al.: The multimodal brain tumor image segmentation benchmark (brats). IEEE Trans. Med. Imaging 34(10), 1993–2024 (2015). https://doi.org/10.1109/TMI.2014.2377694

  15. Ronneberger, O., Fischer, P., Brox, T.: U-net: convolutional networks for biomedical image segmentation. In: Navab, N., Hornegger, J., Wells, W.M., Frangi, A.F. (eds.) MICCAI 2015. LNCS, vol. 9351, pp. 234–241. Springer, Cham (2015). https://doi.org/10.1007/978-3-319-24574-4_28

    CrossRef  Google Scholar 

  16. Sandler, M., Howard, A.G., Zhu, M., Zhmoginov, A., Chen, L.: Mobilenetv 2: inverted residuals and linear bottlenecks. In: 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR, pp. 4510–4520 (2018)

    Google Scholar 

  17. Simi, V., Joseph, J.: Segmentation of glioblastoma multiforme from MR images - a comprehensive review. Egyptian J. Radiol. Nucl. Med. 46(4), 1105–1110 (2015)

    CrossRef  Google Scholar 

  18. Sun, T., Chen, Z., Yang, W., Wang, Y.: Stacked u-nets with multi-output for road extraction. In: 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops (CVPRW), pp. 187–1874 (2018)

    Google Scholar 

  19. Tan, M., et al.: MnasNet: platform-aware neural architecture search for mobile. In: 2019 IEEE Conference on Computer Vision and Pattern Recognition, CVPR, pp. 2820–2828. Computer Vision Foundation/IEEE (2019)

    Google Scholar 

  20. Tan, M., Le, Q.V.: EfficientNet: rethinking model scaling for convolutional neural networks. In: Proceedings of the 36th International Conference on Machine Learning, ICML, vol. 97, pp. 6105–6114 (2019)

    Google Scholar 

  21. Thangarasa, V., Tsai, C.Y., Taylor, G.W., Köster, U.: Reversible fixup networks for memory-efficient training. In: NeurIPS Systems for ML (SysML) Workshop (2019)

    Google Scholar 

  22. Wu, Y., He, K.: Group normalization. In: Proceedings of the European Conference on Computer Vision (ECCV) (2018)

    Google Scholar 

  23. Yao, W., Zeng, Z., Lian, C., Tang, H.: Pixel-wise regression using u-net and its application on pansharpening. Neurocomputing 312, 364–371 (2018)

    CrossRef  Google Scholar 

Download references

Author information

Authors and Affiliations

Authors

Corresponding author

Correspondence to Mihir Pendse .

Editor information

Editors and Affiliations

Rights and permissions

Reprints and Permissions

Copyright information

© 2021 Springer Nature Switzerland AG

About this paper

Verify currency and authenticity via CrossMark

Cite this paper

Pendse, M., Thangarasa, V., Chiley, V., Holmdahl, R., Hestness, J., DeCoste, D. (2021). Memory Efficient 3D U-Net with Reversible Mobile Inverted Bottlenecks for Brain Tumor Segmentation. In: Crimi, A., Bakas, S. (eds) Brainlesion: Glioma, Multiple Sclerosis, Stroke and Traumatic Brain Injuries. BrainLes 2020. Lecture Notes in Computer Science(), vol 12659. Springer, Cham. https://doi.org/10.1007/978-3-030-72087-2_34

Download citation

  • DOI: https://doi.org/10.1007/978-3-030-72087-2_34

  • Published:

  • Publisher Name: Springer, Cham

  • Print ISBN: 978-3-030-72086-5

  • Online ISBN: 978-3-030-72087-2

  • eBook Packages: Computer ScienceComputer Science (R0)