UNETR++: Delving into Efficient and Accurate 3D Medical Image Segmentation

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Shaker, Abdelrahman, Maaz, Muhammad, Rasheed, Hanoona, Khan, Salman, Yang, Ming-Hsuan, Khan, Fahad Shahbaz
Formato: Preprint
Publicado: 2022
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866916234698686464
author Shaker, Abdelrahman
Maaz, Muhammad
Rasheed, Hanoona
Khan, Salman
Yang, Ming-Hsuan
Khan, Fahad Shahbaz
author_facet Shaker, Abdelrahman
Maaz, Muhammad
Rasheed, Hanoona
Khan, Salman
Yang, Ming-Hsuan
Khan, Fahad Shahbaz
contents Owing to the success of transformer models, recent works study their applicability in 3D medical segmentation tasks. Within the transformer models, the self-attention mechanism is one of the main building blocks that strives to capture long-range dependencies. However, the self-attention operation has quadratic complexity which proves to be a computational bottleneck, especially in volumetric medical imaging, where the inputs are 3D with numerous slices. In this paper, we propose a 3D medical image segmentation approach, named UNETR++, that offers both high-quality segmentation masks as well as efficiency in terms of parameters, compute cost, and inference speed. The core of our design is the introduction of a novel efficient paired attention (EPA) block that efficiently learns spatial and channel-wise discriminative features using a pair of inter-dependent branches based on spatial and channel attention. Our spatial attention formulation is efficient having linear complexity with respect to the input sequence length. To enable communication between spatial and channel-focused branches, we share the weights of query and key mapping functions that provide a complimentary benefit (paired attention), while also reducing the overall network parameters. Our extensive evaluations on five benchmarks, Synapse, BTCV, ACDC, BRaTs, and Decathlon-Lung, reveal the effectiveness of our contributions in terms of both efficiency and accuracy. On Synapse, our UNETR++ sets a new state-of-the-art with a Dice Score of 87.2%, while being significantly efficient with a reduction of over 71% in terms of both parameters and FLOPs, compared to the best method in the literature. Code: https://github.com/Amshaker/unetr_plus_plus.
format Preprint
id arxiv_https___arxiv_org_abs_2212_04497
institution arXiv
publishDate 2022
record_format arxiv
spellingShingle UNETR++: Delving into Efficient and Accurate 3D Medical Image Segmentation
Shaker, Abdelrahman
Maaz, Muhammad
Rasheed, Hanoona
Khan, Salman
Yang, Ming-Hsuan
Khan, Fahad Shahbaz
Computer Vision and Pattern Recognition
Owing to the success of transformer models, recent works study their applicability in 3D medical segmentation tasks. Within the transformer models, the self-attention mechanism is one of the main building blocks that strives to capture long-range dependencies. However, the self-attention operation has quadratic complexity which proves to be a computational bottleneck, especially in volumetric medical imaging, where the inputs are 3D with numerous slices. In this paper, we propose a 3D medical image segmentation approach, named UNETR++, that offers both high-quality segmentation masks as well as efficiency in terms of parameters, compute cost, and inference speed. The core of our design is the introduction of a novel efficient paired attention (EPA) block that efficiently learns spatial and channel-wise discriminative features using a pair of inter-dependent branches based on spatial and channel attention. Our spatial attention formulation is efficient having linear complexity with respect to the input sequence length. To enable communication between spatial and channel-focused branches, we share the weights of query and key mapping functions that provide a complimentary benefit (paired attention), while also reducing the overall network parameters. Our extensive evaluations on five benchmarks, Synapse, BTCV, ACDC, BRaTs, and Decathlon-Lung, reveal the effectiveness of our contributions in terms of both efficiency and accuracy. On Synapse, our UNETR++ sets a new state-of-the-art with a Dice Score of 87.2%, while being significantly efficient with a reduction of over 71% in terms of both parameters and FLOPs, compared to the best method in the literature. Code: https://github.com/Amshaker/unetr_plus_plus.
title UNETR++: Delving into Efficient and Accurate 3D Medical Image Segmentation
topic Computer Vision and Pattern Recognition
url https://arxiv.org/abs/2212.04497