Self-attention on Multi-Shifted Windows for Scene Segmentation

The scene segmentation model is built on the recently proposed Swin Transformer Liu et al. (2021b), which is a powerful visual Transformer in large-scale image classification and can serve as a general-purpose backbone for down-stream learning tasks in computer vision.

Swin Transformer constructs hierarchical feature maps just like the commonly used Convolution Network (ConvNet) structures such as ResNet He et al. (2016) and DenseNet Huang et al. (2017). However, ConvNets and Swin Transformer have fundamental differences in computing feature representations. ConvNets apply convolution operators to overlapped receptive fields to progressively learn local semantic information, while Swin Transformer learns the spatial feature dependencies within the shifted image windows. The ability of learning feature dependency in the shifted image windows endows Transformer model a compelling advantage over FCNs for scene segmentation, since it explicitly describes the spatial dependencies. Another important difference is the resolution reduction between the learning stages. ConvNets use either doubled strides or down-samplings, while Swin Transformer adopts the patch-merging scheme. The above two differences make it inappropriate to directly build an encoder with Swin Transformer just like FCNs, i.e., keeping the feature resolution in the deeper stages and increasing the dilation rate with atrous convolution to enlarge the receptive field Zhao et al. (2017); Fu et al. (2019a). The reasons are illustrated as follows:

• •

  The learned embedding of merged patches no longer provides the accurate representation for the down-streamed computations, if keeping the high-resolution features in the deeper stages;

• •

  If we calculate the dependency of non-adjacent patches like the way of atrous convolution Zhao et al. (2017); Chen et al. (2018), the self-attentions still learn the same dependency within a window.

Based on the above observations, in the design of the encoder backend on Swin Transformer, the structure of Feature Pyramid Network (FPN) with the Pyramid Pooling Module (PPM) is used as a whole learning framework (UperNet) for scene segmentation Liu et al. (2021b). However, the PPM is applied on the low-resolution but high-dimensional feature maps in UperNet, which is computationally expensive and ineffective in restoring the spatial details. In our work, we only use the structure of Feature Pyramid Network (FPN), which can represent hierarchical features and carry rich spatial-semantic information, but remove the PPM and add a few Swin blocks before upsampling. In deep neural networks, FPN has been successfully used for object detection Lin et al. (2017) to extract multi-scale feature maps. The Transformer-based FPN is illustrated in Figure 2(a), which composes of both bottom-up and top-down pathways. The bottom-up pathway is the four-stage computation streamline in Swin Transformer. As we go up, the spatial resolution decreases due to the patch merging. The semantic value for each stage increases with more high-level feature representation is learned. The feature outputs from the bottom layers are in high resolution but the semantic value is limited. By contrast, the features from the top layers are in low-resolution but with rich semantics. In scene segmentation, the features with high-resolution and rich semantics are both necessary for pixel-level classification, so FPN provides a top-down pathway to construct higher resolution layers and keep the rich semantic information.

Assume an input image is $\mathbf{X}_{0}\in\mathbb{R}^{H\times W\times 3}$, where $H$ and $W$ are the width and height, respectively. In bottom-up pathway in Swin Transformer, the output features of the four stages are $\mathbf{X}_{1}\in\mathbb{R}^{\frac{H}{4}\times\frac{W}{4}\times C},\mathbf{X}_{2}\in\mathbb{R}^{\frac{H}{8}\times\frac{W}{8}\times 2C},\mathbf{X}_{3}\in\mathbb{R}^{\frac{H}{16}\times\frac{W}{16}\times 4C}$ and $\mathbf{X}_{4}\in\mathbb{R}^{\frac{H}{32}\times\frac{W}{32}\times 8C}$, respectively, where $C$ is a constant depending on the model size of the backbone. In the top-down pathway, it progressively hallucinates higher resolution features by upsampling spatially coarser but semantically richer feature maps from higher pyramid levels. These features are then fused with the feature maps computed from the bottom-up pathway via lateral connections. The lateral outputs $\mathbf{L}_{1},\mathbf{L}_{2},\mathbf{L}_{3}$ and $\mathbf{L}_{4}$ in the top-down pathway are computed as follows:

where $\text{UP}_{\times 2}(\cdot)$ is the upsampling operator with bilinear interpolation by 2, and $\text{LP}(\cdot)$ is the linear projection with a batch normalization followed by a ReLU activation. After that, we apply four window-based self-attention blocks on the lateral outputs, then sum them up to form the final feature output $\mathbf{Y}_{0}$ of the encoder, as follows:

where W-MSA is the Window based Multi-head Self-attention module. Different from the use of FPN for object detection in Lin et al. (2017) that predicts multi-scale object sizes with pyramid feature maps, we aggregate the multi-scale lateral features in the same feature resolution. The resolution of the final feature output $\mathbf{Y}_{0}$ is $\frac{H}{4}\times\frac{W}{4}$.

The encoder with FPN is similar to the Unified Perceptual Parsing Network (UPerNet) proposed in Xiao et al. (2018). However, in our design we remove the pyramid pooling module (PPM)Zhao et al. (2017) and apply the Swin blocks to aggregate the feature maps, thus the proposed encoder does not contain any convolution operator. Note that the PPM needs more memory and extra $3\times 3$ convolutions, which is computationally expensive. In the experiment, we prove that even with the self-attention on multi-shifted windows, the whole learning framework of MSwin is more computationally efficient than UPerNet. Using FPN as the encoder of the scene segmentation model can well suit the hierarchical backbone of Swin Transformer without hurting the internal computation structure. Furthermore, the feature output carries very rich semantic information and keeps the high resolution of features. So if we directly add a classification layer as the decoder on the top of the FPN encoder output, it can form a pure Vision Transformer model for scene segmentation. Here we name it Transformer based FPN (T-FPN), which serves as a baseline in the experiment. Note that the T-FPN only contains the self-attention on single-shifted windows.

We next introduce how to design the decoder by exploring the spatial dependencies with self-attention on multi-shifted windows.

The Shifted Window based Multi-head Self-attention (SW-MSA) is a simple yet effective module in Swin Transformer to enrich the feature representation. The intuition behind SW-MSA is to introduce cross-window connections of non-overlapping windows thus improving the modelling power of W-MSA. In scene segmentation tasks, it is commonly recognized that learning on multi-scale feature representations is beneficial to take different contextual information into consideration, thus giving more accurate predictions. In our model, we also follow this strategy in the design of segmentation head by applying self-attention on multi-shifted windows.

Here we denote $m$ and $n$ be the window size and the shifted size in a SW-MSA module, where $n<m$. W-MSA is a special case of SW-MSA when $n=0$. Note that the self-attention is learned to describe the spatial dependency among non-overlapping windows, so different settings of $m$ and $n$ essentially implement multiple self-attentions in the same image. In Figure 1 we illustrate two shifted windows, where an image feature map with the resolution $6\times 6$ can be partitioned to either $2\times 2$ or $3\times 3$ windows. Applying two shifted sizes 2 and 1 on these partitions, we can obtain four different outputs by setting $m=3,n=0$, $m=3,n=2$, $m=2,n=0$, and $m=2,n=1$ in the SW-MSA modules, respectively. This can implement a multi-scale attention within different sized windows thus can diversify the feature map representation to improve the discriminative power of self-attention for pixel-wise classification. For simplicity, we set $n=\lfloor\frac{m}{2}\rfloor$ in the self-attention on multi-shifted windows of the decoder. Assume we have $L$ SW-MSA modules, and the feature output from the T-FPN encoder is $\mathbf{Y}_{0}$, the three structures for the decoder are described as follows.

PASCAL VOC 2012Everingham et al. (2015) contains 20 foreground object classes and one background class. The original dataset has 1,464 and 1,449 images for training and validation, respectively. To augment the training dataset, we also use extra annotations provided by Hariharan et al. (2011), so the total number of training images is 10,582. Here we do not use MS COCO dataset to pretrain the segmentation model.

COCO-Stuff 10KCaesar et al. (2018) is a subset of the complete COCO-Stuff dataset, which provides pixel-wise semantic labels for the whole scene, including both “thing” and “stuff” classes. It contains 9,000 and 1,000 images for training and validation (testing), respectively. Following Ding et al. (2018); Fu et al. (2019a); Yuan et al. (2020), we evaluate the segmentation performance on 171 categories (80 objects and 91 stuff) to each pixel.

ADE20KZhou et al. (2017) is a challenging scene parsing benchmark, in which the images are from both indoor and outdoor environments and are annotated by 150 fine-grained semantic concepts. It contains 20,210, 2,000 and 3,352 images for training, validation and testing, respectively.

CityscapesCordts et al. (2016) is an urban traffic dataset, in which the images are densely annotated by 19 classes. We train the scene segmentation model with the finely annotated 2,975 images for training and 500 images for validation, respectively. It also provides 20,000 coarsely labeled images to pre-train the segmentation models.

Our experimentation is based on the semantic segmentation package mmsegmentation GIT (2020). The MSwin models implemented in the experiment are backend on a small-sized Swin-S and a medium-sized Swin-B Liu et al. (2021b), which are pretrained on ImageNet-1K and ImageNet-22K, respectively. The window sizes in Swin-S and Swin-B are $7\times 7$ and $12\times 12$. In all decoder heads that contain SW-MSA modules, we set the dimension of embedding to 512 and 8-heads attention without any change. For the attention on multi-shifted windows, we used three different window sizes $5\times 5$, $7\times 7$ and $12\times 12$, and the shifted sizes were set to $2\times 2$, $3\times 3$ and $6\times 6$, respectively. Considering the combinations of both W-MSA and SW-MSA, there are totally $L=6$ attention blocks. Following Zhao et al. (2017); Fu et al. (2019a), we added an auxiliary loss head composed of a two-layer sub-network. The auxiliary loss and main loss were computed concurrently with weights 0.4 and 1, respectively. We used AdamW optimizer L. and H. (2019) with the initial learning rate $6\times 10^{-5}$ and weight decay 0.01 after each iteration. For image augmentation, we applied random cropping, random flipping and photometric distortion. On the Cityscapes dataset, we applied random cropping with the size $512\times 1024$. On the COCO-Stuff 10K, we set a smaller cropping size $480\times 480$. On the rest two datasets, the cropping sizes were set to $512\times 512$. We used categorical cross-entropy as loss function and report the mean Intersection-over-Union (mIoU) score in the evaluation of scene segmentation performance. These are consistent with the optimization settings of other baselines for fair comparisons, although applying some recently proposed learning objectives such as Lovász-softmax Berman et al. (2018) and Margin calibrated log-loss Yu et al. (2021) can further improve the mIoU scores. We used the mixed-precision and gradient checkpoint in the model training, which allows us to set a larger mini-batch size and can effectively save the GPU memory usage without hurting the normalization layers. Our experiments were conducted on a server equipped with two NVIDIA Tesla V100 GPU cards.

To demonstrate the effectiveness of attention on multi-shifted windows, we conducted an ablation study of different window-size combinations on top of the T-FPN model backend on Swin-S. We separately set the window sizes $5\times 5$, $7\times 7$ and $12\times 12$, as well as their combinations. We tested the proposed MSwin-P and MSwin-S, and the single-scaled mIoU predictions are illustrated in Table 1. We can observe that it is unclear which single window size obtains the best results, and applying the self-attention with two window sizes do not necessarily improve the model performance. However, when we use three different window sizes by setting $L=6$, the two MSwin models perform slightly better and have relatively stable results. It is expected that using more self-attention in MSwin model with more window sizes can further boost the segmentation performance, but the computational complexity also increases accordingly.

We then experimented the segmentation models with Swin-S and Swin-B as backbones to verify the effectiveness of the baseline T-FPN, as well as three MSwin models on this dataset. The mIoU scores on the validation set are summarized in Table 2, where SS and MS are the abbreviates of single-scale and multi-scale prediction for mIoU scores, respectively. From the table, we can see that even with a softmax classifier as a decoder, the T-FPN backed on Swin Transformer is a powerful semantic segmentation model. Applying the small-sized backbone Swin-S, the proposed MSwin-P, MSwin-S and MSwin-C further improve the baseline T-FPN by 0.89%, 1.28% and 0.71%, and by 0.88%, 0.67% and 0.63% in terms of mIoU score when using single-scale and multi-scale predictions, respectively. When applying the more powerful backbone Swin-B, the three MSwin models further improve the performance of the baseline similarly. Figure 3 gives some visualization results of the four segmentation methods on the validation dataset.

We mixed the train+val datasets, fine-tuned the models, and used the multi-scale prediction on the test set, then uploaded the results to the evaluation server. In Table 3, we compared our models with the recently proposed methods. The three MSwin models achieve the new state-of-the-arts without the pre-training on MS COCO dataset.

We conducted experiments on the COCO-Stuff 10K dataset to prove the generalization ability of the proposed methods. All models were optimized by 40,000 iterations. The comparisons with previously state-of-the-arts are reported in Table 4. Our pure Transformer-based models, including the baseline T-FPN, consistently outperform all the ConvNet-based counterparts. When applying three different aggregation strategies in the decoder and multi-scale prediction, MSwin-P, MSwin-S and MSwin-C further boost the baseline T-FPN by 0.3%, 1.0% and 0.5% in terms of mIoU score, respectively. Our models achieve a slightly lower mIoU compared to the OCRNet backend on HRFormer-B on this dataset.

On this dataset, we used both Swin-S and Swin-B as backbones for model training to evaluate the performance. Table 5 shows the comparisons of FLOPs, as well as the results on the validation set. We can observe that the Transformer-based segmentation models generally needs more computational resources, except the recently proposed SegFormer. However, backend with Swin Transformer, our models are more computationally efficient. Applying the self-attention on multi-shifted windows in the decoder doubles the FLOPs in MSwin compared to T-FPN, but leads to more accurate segmentation results. On the validation dataset, all Transformer-based segmentation models substantially outperform the ConvNet counterparts by a large margin. The reason for this is mainly because the self-attention mechanism in Transformers has a very strong capability in modelling the spatial dependency, which can effectively parse very complex scenes in different visual environments by capturing the long-range dependencies. On the validation set, our MSwin models achieve comparable results to UPerNet backend on Swin Transformer, SETR backend on ViT-L and SegFormer. Some segmentation examples are illustrated in Figure 4.

We then fine-tuned the MSwin-P model on the train+validation set to improve the overall classification accuracy. On the evaluation server, MSwin-P obtains the pixel-wise accuracy of 0.767, mIoU 0.457 and the final test score of 0.612.

On this dataset, we directly used Swin-B as the backbone to evaluate the effectiveness of the proposed MSwin models in street scene segmentation. Only using the fine-labeled dataset, we show the improvements brought by the decoder with self-attention on multi-shifted windows based on the T-FPN encoder in Table 6. On the validation set, the parallel structure of MSwin achieves the best mIoU score when using single-scale prediction, which improves mIoU by 0.67% compared to T-FPN. Applying the sequential structure, the model obtains the best performance with multi-scale prediction, which outperforms the baseline by 0.62% of mIoU. We show some segmentation examples and the confusion matrices of dense prediction in Figure 5 and 6, respectively. Compared with a thin decoder based on T-FPN, self-attention on multi-shifted windows can further reduce false positives and lead to finer details of street scenes.

We pre-trained MSwin-S and MSwin-C on the coarsely labeled training data then fine-tuned them on the fine labeled training dataset. After that, we applied the multi-scale prediction on the test set and submitted the results to the evaluation server. The overall comparisons of our models with some recently proposed methods are summarized in Table 7. We can see that without the pre-training with the coarsely labeled images, the proposed MSwin-P obtains the mIoU score 81.8%, which is slightly worse than the best scene segmentation models based on ConvNets. When pre-trained with the coarsely labeled data, MSwin-S and MSwin-C obtain the best results. Specifically, compared to the Transformer-based scene segmentation models SETR and SegFormer, MSwin-S and MSwin-C outperform the second-best by 0.5% and 0.3% in terms of mIoU score, respectively.