nnUnetFormer: an automatic method based on nnUnet and transformer for brain tumor segmentation with multimodal MR images

Overall, our nnUnetFormer consists of a modified nnUnet as the main network framework and the embedded transformer modules, as illustrated in figure 1. It takes the multimodal MR images of patch size 128 × 128 × 128 as input to obtain the fused local and global features, and finally outputs the predicted segmentation maps. We first utilize the encoder of our nnUnetFormer to efficiently generate a series of feature maps by extracting multi-scale spatial features. In this way, rich local 3D context features from the encoder are effectively embedded in high-level features step by step, which are subsequently fed into the transformer modules to further learn the long-range dependencies and supplement the local context features with a global receptive field. After that, we repeatedly apply the concatenation operation on the corresponding two branches of up-sampling and skip-connection, followed by convolutional layers to gradually produce a high-resolution segmentation map.

Zoom In Zoom Out Reset image size

Concretely, since the features in shallower layers tend to generate more detailed spatial information and those in deeper layers extract more discriminative semantic information (Jia et al 2022), we have modified the original nnUnet architecture by using a deeper encoder, which was inspired by Luu and Park (2022). This modification aims to increase the network capacity to better model greater data diversity. Here, we used 6 levels of same-resolution convolutional layers and 5 strided convolution (stride = 2) operations for down-sampling in the encoder to extract more discriminative semantic information. The decoder followed the same structure with a concatenation operation and two convolution operations. Note that the concatenation operation is used to concatenate the transpose convolution (TransConv) operations for up-sampling with the skip features from the encoder branch at the same-resolution level. To be more specific, the combination of batch normalization and Leaky ReLU (LReLU) with negative slope of 0.01 was applied after every convolution operation; the number of channels was set from 4 to 64 at the initial level (level 1) and doubled for every down-sampling. Besides, with respect to the initial nnUnet architecture, additional sigmoid outputs were added to every resolution except for the two lowest levels as auxiliary supervision branches according to deep supervision learning technique (Wang et al 2015) for enhancing gradient propagation to the earlier layers as well as improving the speed of deeper network convergence in training phase.

We found that due to the GPU memory constraint, it was inadvisable to apply the transformer module to the high-resolution features. Hence, we fused up to three transformer modules into the three lowest resolution scales (Level 4, Level 5, and Level 6) of our modified nnUnet while keeping the skip connections unchanged. The structure of the transformer module is illustrated in figure 2. Specifically, assuming the input of our transformer layer is the feature map $F\in {R}^{C\times W\times H\times D},$ which firstly goes through a convolution operation to change its channel dimension from C to the preassigned number n, and then reshapes the 3D spatial dimensions of the new feature representations after the convolution operation as 1D dimension of size S (numerically equal to W × H × D). This process is named linear projection, resulting in feature map $F\text{'}\in {R}^{n\times S}.$ After the linear projection, we add a 1D learnable positional embedding (PE) module to $F\text{'},$ forming S n-dimensional embedded patches (or tokens). Subsequently, a stack of transformer layers, which mainly consists of multi-head attention (MHA) and multilayer perceptron (MLP) sub-layers, was used to learn long-range dependencies with global receptive field. Let F_l denote the output of the lth (l $\in$ [1, 2, ..., L]) transformer layer, it can be calculated by:

$$\begin{eqnarray}&&{F}_{l-1}^{\text{'}}=\mathrm{MHA}\left(\mathrm{LN}\left({F}_{l-1}\right)\right)+{F}_{l-1}\end{eqnarray} \tag{ 1 }$$

$$\begin{eqnarray}&&{F}_{l}=\mathrm{MLP}\left(\mathrm{LN}\left({F}_{l-1}^{\text{'}}\right)\right)+{F}_{l-1}^{\text{'}},\end{eqnarray} \tag{ 2 }$$

where LN(*) is the layer normalization. In transformer module, the core element is a MHA sub-layer, which was used to update the states of each embedded patch by aggregating information globally. The MHA sub-layer comprises P parallel self-attention (SA) heads and one linear layer. Each head has its own learnable weight matrices. The SA block is a parameterized function that learns the mapping between a query and the corresponding key and value representations in a feature Fa (Hatamizadeh et al 2022). Let Q, K, and V denote respectively query, key and value; the output of ith SA head is defined as

$$\begin{eqnarray}&&{\rm{SA}}\left(F{a}_{i}\right)={\rm{softmax}}\left(\displaystyle \frac{{Q}_{i}{K}_{i}^{T}}{\sqrt{d}}\right){V}_{i}\end{eqnarray} \tag{ 3 }$$

where i $\in$ [1, 2, ..., P], and d is the dimension of Q and K. The MLP sub-layer is composed of two linear layers linked by one GELU activation function. To fit the input dimension of 3D CNN decoder, we use a feature mapping module to project the patches back to the feature map, which owns the same dimension as the input of the transformer module.

Zoom In Zoom Out Reset image size

In MR images of brain tumor patients, the volume of whole tumor (WT) region is much smaller than that of the whole brain, while either TC or ET region occupies even less, which makes the BTS task suffer from a serious class imbalance problem. Dice loss (Milletari et al 2016) is based on the Dice coefficient similarity, which is a set-based metric for measuring the overlap between segmentation results and annotation labels in semantic segmentation. To balance the loss between imbalanced classes, weighted cross-entropy (WCE) (Akil et al 2020) loss applies a weighting factor to the cross-entropy loss for increasing the penalty to the minority classes over the majorities. Considering both Dice loss and WCE loss can be used to alleviate the class imbalance problem, we adopt a combination of them as the training loss for our network, each of which is expressed as follows:

$$\begin{eqnarray}&&{L}_{\mathrm{Dice}}=1-2\frac{\displaystyle {\sum }_{j=0}^{M}{g}_{j}{p}_{j}}{\displaystyle {\sum }_{j=0}^{M}\left({g}_{j}+{p}_{j}\right)}\end{eqnarray} \tag{ 4 }$$

$$\begin{eqnarray}&&{L}_{\mathrm{WCE}}=-\displaystyle \sum _{j=0}^{M}{w}_{j}{g}_{j}\,\mathrm{log}\left({p}_{j}\right)\end{eqnarray} \tag{ 5 }$$

$$\begin{eqnarray}&&{L}_{\mathrm{seg}}={L}_{\mathrm{Dice}}+{L}_{\mathrm{WCE}}\end{eqnarray} \tag{ 6 }$$

where M is the number of classes, g_j the ground-truth for class j, p_j the model prediction for class j, and w_j the weight of class j.

In view of sigmoid outputs we added (figure 1), our network was trained to optimize the combined loss, calculated at the final full resolution output as well as at the auxiliary outputs of lower resolution. The combined loss L_combined is defined as

$$\begin{eqnarray}&&{L}_{\mathrm{combined}}=\displaystyle \sum _{o=0}^{O}{\mathrm{weight}}_{o}\times {L}_{\mathrm{seg}}\left(o\right)\end{eqnarray} \tag{ 7 }$$

where O represents the number of the outputs at various resolutions, o the o-th output of the model, L_seg(o) the loss of the o-th output. Especially, the weights for each output from full resolution (o = 0) to lower resolution are calculated as follows:

$$\begin{eqnarray}&&{\mathrm{weight}}_{o}={2}^{-o}/\displaystyle \sum _{o=0}^{O}{2}^{-o}.\end{eqnarray} \tag{ 8 }$$

The multimodal magnetic resonance images used for the model training and evaluation in this study were obtained from the publicly available datasets of BraTS 2021 and Federated Tumor Segmentation challenge (FeTS) 2021 (Isik-Polat et al 2022). In BraTS 2021 and FeTS 2021, their training sets contain 1251 and 341 cases collected from multiple institutions, respectively. In FeTS 2021, the cases in the training set were partitioned into 22 clients according to their original institutions and tumor sizes (Yin et al 2022), and the statistical results of the cases in each client is illustrated in figure 3. In the training set of both datasets, each case contains four 3D modalities (T1, T2, T1ce and Flair) MR images along with the corresponding segmentation annotations of tumor sub-types. All the MR images and the associated annotations were resized to 240 × 240 × 155 and stored as NIfTI files (.nii.gz). Following the two challenges, all the MR images have been preliminarily pre-processed by organizers: co-registered to the same anatomical template, resampled to 1 mm³ isotropic resolution, and skull-stripped. The annotations of each case have been achieved manually following the same annotation protocol and comprise the GD-enhancing tumor (label 4), the peritumoral edema (ED, label 2), the necrotic and non-enhancing tumor (NCR/NET, label 1), and the background (label 0). Note that label 3 does not appear here since it was merged into label 1 by the organizers. Furthermore, the almost homogeneous annotations have been clustered together to compose three mutually inclusive tumor sub-regions, representing the enhancing tumor (ET: label 4), tumor core (TC: label 1 + label 4), and whole tumor (WT: label 1 + label 2 + label 4), respectively.

Zoom In Zoom Out Reset image size

As for BraTS 2021, to overcome the bias caused by a particular selection of the pair of training and testing sets, each network was trained from scratch and evaluated using 5-fold cross-validation. Hence, the performance evaluated on this dataset was measured using the averaged evaluation metrics from the 5 validation folds. Moreover, since the purpose of task 2 in FeTS 2021 challenge was to evaluate the robustness and generalization of different segmentation algorithms (Nalawade et al 2022, Yin et al 2022), we further used FeTS 2021 to conduct more experiments to evaluate the generalization capability of our method. Here, following the principle of making up the training set with as few the institutions or centers as possible, we used 257 cases from clients 11, 13, 14, 15, 16, 17, 18, 19, 20, 21 and 22 for training and the remaining 84 cases from other unseen 11 clients (independent institutions, to be precise) for testing with an approximate 3:1 ratio.

To effectively evaluate the performance of the proposed nnUnetFormer method, the official evaluation metrics of Dice similarity coefficient (DSC) and 95th percentile of Hausdorff distance (HD95) were used in this study. DSC was used to measure the spatial overlap between the segmentation predictions and the ground-truth annotations, which is numerically equal to 1−L_Dice. Hausdorff distance (HD) originates from set theory and measures the maximum distance of a point set to the nearest point in another set. As opposed to DSC, HD represents a global measure for spatial overlap and is sensitive to local differences (Momin et al 2022). DSC and HD are defined by:

$$\begin{eqnarray}&&\mathrm{DSC}=\frac{2| G\cap P| }{| G| +| P| }\end{eqnarray} \tag{ 9 }$$

$$\begin{eqnarray}&&\mathrm{HD}\left(G,P\right)=\,\max \left\{\mathop{\max }\limits_{g\in G}\mathop{\min }\limits_{p\in P}| | g-p| | ,\mathop{\max }\limits_{p\in P}\mathop{\min }\limits_{g\in G}| | p-g| | \right\},\end{eqnarray} \tag{ 10 }$$

where G denotes the ground-truth annotations, and P the segmentation predictions. For each metric, three tumor regions of ET, TC and WT were evaluated individually, and we took the average DSC of the three regions as the main metric to evaluate the performances of trained models.

Our proposed nnUnetFormer method is implemented using PyTorch. Keeping consistent with the training methodology of nnUnet, we also adopted the region-based strategy for training instead of predicting the mutually exclusive brain tumor sub-types, and used the cropping and z-score normalization operations for original MR images in the pre-processing step. Considering the over-fitting phenomenon is almost unavoidable in deep learning methods when training a large neural network containing a huge number of parameters to learn with limited training data (Liu et al 2021a, Zhang et al 2021b), we applied a large variety of spatial data augmentation techniques, including random rotation, random cropping, random scaling, random elastic deformation, and random gamma correction on the fly during network training to boost the generalization capability of our model. In the training phase, the initial learning rate was set to 1 × 10⁻² with the updating strategy of scheduler optimization. The optimizer used stochastic gradient descent algorithm with weight decay of 5 × 10⁻⁵ and Nesterov momentum of 0.99. The batch size was set to 2. The number of training epoch was set to 200. The number of SA heads in transformer layer was set to 8.

In this study, we took the original nnUnet as baseline. Only one transformer module was embedded into level 4, level 5, level 6, thus leading to nnUnetFormer (Level 4), nnUnetFormer (Level 5), and nnUnetFormer (Level 6), respectively. When embedding one transformer module in each of levels 4, 5, 6, we have nnUnetFormer (Levels 4–6). To validate the performance under different modifications on baseline, a variety of ablation studies were performed. Paired Wilcoxon signed-rank tests were conducted for the results between the nnUnet baseline and our modified models at the significance level of 0.05. Firstly, the quantitative validation results are given in table 1. We can observe that our nnUnetFormer (Levels 4–6) model achieves best overall performance both with DSC 0.936, 0.921 and 0.872, and with HD95 3.96, 4.57 and 10.45 for the three regions of WT, TC, and ET, respectively. Especially, the overall performance of our nnUnetFormer (Levels 4–6) model was slightly better than baseline, with the average DSC increased by 1.22% (0.910 versus 0.899). Besides, incorporating the modifications into baseline led to a consistent improvement in terms of average DSC. In particular, integrating the three transformer modules into our modified nnUnet architecture (model D) improved the performance for the task of BTS, with the average DSC increased by 0.55% (0.910 versus 0.905). As for our nnUnetFormer (Level 4), nnUnetFormer (Level 5), and nnUnetFormer (Level 6) models, we found that no matter the transformer module was embedded into Level 4, Level 5 or Level 6, we achieved better segmentation performance in terms of both average DSC and average HD95 compared with other models without using transformer module. By comparing the runtime (FLOPs) and the number of learning parameters of each model, it can be seen that our nnUnetFormer model has no advantages. Secondly, for better understanding our nnUnetFormer method, we also illustrate an example of segmentation results using different modifications based on nnUnet baseline on BraTS 2021, as shown in figure 4. From table 1 and figure 4, we can observe that the segmentation results are gradually improved when the proposed modifications are integrated, especially for the TC and ET regions.

Zoom In Zoom Out Reset image size

We quantitatively compared our proposed nnUnetFormer method in our internal 5-fold cross-validation split against other recent state-of-the-art methods such as HNF-Netv2 (Jia et al 2022), ResUNet (Pei and Liu 2022), nnSegNet (Jabareen and Lukassen 2022), BiTr-Unet (Jia and Shu 2022), E1D3 Unet (Bukhari and Mohy-ud-Din 2022), SGEResU-Net (Liu et al 2022), MMEF-nnUNet (Huang et al 2022), VT-UNet (Peiris et al 2021), SegResNet (Rahman Siddiquee and Myronenko 2022), Extending nnUnet (Luu and Park 2022) on BraTS 2021 dataset. Considering we have reproduced three networks VT-UNet, SegResNet, and Extending nnUnet according to the public available codes, paired Wilcoxon rank-sum tests for the results were conducted between them and our nnUnetFormer method. From the quantitative experimental results summarized in table 2, we found that our nnUnetFormer (Levels 4–6) model shows slightly better in terms of average DSC and average HD95 compared with the other methods and outperforms the second best method by 0.66% (0.910 versus 0.904) and 6.92% (6.32 versus 6.79), respectively. Specifically, our model also achieved the best DSC scores in WT and TC regions, although its DSC score of ET region is slightly lower than MMEF-nnUNet method. As for TC region, our model outperformed the second best method by 1.1% (0.921 versus 0.911) in terms of DSC score.

Furthermore, we also conducted qualitative analysis by visually comparing the segmentation results of different methods including VT-UNet, SegResNet, Extending nnUnet and our nnUnetFormer (Levels 4–6). From figure 5, we observed that different methods have different segmentation results for the three tumor regions, and our segmentation results are relatively close to the ground-truth compared with the other three methods. Specifically, row 1 indicates that our model can segment the peritumoral edema relatively well while the GD-enhancing tumor and the NCR/NET are accurately segmented. Rows 2 and 3 indicate that our model can capture and identify small targets (NCR/NET and GD-enhancing tumor) better. Note that as for the MR images of patient in row 3, the number of voxels in ET region is only 94, which is much smaller than the number of voxels in TC region.

Zoom In Zoom Out Reset image size

Taking the segmentation performance, training/validation time, and parameters number of our four different nnUnetFormer models into consideration, we used the nnUnetFormer (Level 6) model to conduct more experiments on FeTS 2021 for evaluating the generalization capacity of our nnUnetFormer method. The results of different unseen clients in the testing set are summarized in table 3. Notably, although the segmentation performance of our nnUnetFormer model varies from client to client, it meets the requirements of clinical applications for most clients. What is worth mentioning is that there is a significant difference in image quality between the 3 clients (clients 9, 10, 12), thereby a depressing performance drop when our model shifted from the training set to the unseen cases coming from these 3 independent institutions. Especially, for the cases without GD-enhancing tumor, there are 3.5% in our training set, while there are respectively 44.4%, 66.7%, 66.7% in clients 9, 10, 12 and 0.0% in the other 8 unseen clients. Taking all the testing cases into consideration, our nnUnetFormer model achieves overall segmentation performance with DSC 0.912, 0.872 and 0.759, and HD95 6.16, 8.81 and 38.50 for the regions of WT, TC, and ET, respectively. Besides, the overall DSC performance for all cases in the testing set was also demonstrated with box and whisker plots in figure 6 for better illustration. From figure 6, we can easily observe that our nnUnetFormer method achieves high-performance segmentation predictions for WT, TC and ET regions in almost all patients. Finally, we also compared our method with the methods we reproduced such as VT-UNet (Peiris et al 2021), SegResNet (Rahman Siddiquee and Myronenko 2022), Extending nnUnet (Luu and Park 2022) on FeTS 2021, as shown in table 4. From table 4, we can find our results are slightly better for the segmentation with best average DSC and average HD95. The average DSC score of our method exceeds the methods of VT-UNet, SegResNet, and Extending nnUnet by 1.4% (0.848 versus 0.836), 2.4% (0.848 versus 0.828), and 0.6% (0.848 versus 0.843), respectively. Furthermore, more quantitative performance comparison with five state-of-the-art methods using DSC score collected from corresponding papers is shown in figure 7. Overall, although our method did not pre-process the images coming from different clients, it is better than the other methods in segmenting both WT and TC regions, except that the DSC score of ET region is slightly lower than that of Yin et al by 0.1%, which was implemented by training a nnUnet model with all cases in FeTS 2021 training set and took the first place for task 2. Particularly, as for TC region, the DSC score of our method exceeds the methods of Nalawade et al (2022), Mächler et al (2022), Isik-Polat et al (2022), Khan et al (2022), and Yin et al (2022) by 12.8% (0.872 versus 0.773), 10.0% (0.872 versus 0.793), 17.7% (0.872 versus 0.741), 22.1% (0.872 versus 0.714), and 12.8% (0.872 versus 0.773), respectively.

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size