Cross-scale sampling transformer for semantic image segmentation

2.1 Boundary handling

Much of the previous work focuses on exploring semantic boundaries [40, 60]. These methods are often applied to high-resolution images but are ineffective for context modeling and multi-scale fusion. In the case of many categories, it is easy to lead to classification errors. Seman-flow [31] applies optical flow to semantic segmentation and proposes a module of pixel deviation alignment to calibrate the offset pixels. FaPN [25] uses the spatial position alignment module to calibrate the pixel position deviation caused in the upsampling process. Fekri et al. [17] proposes a method for bark texture classification with high accuracy based on the improved local ternary patterns (ILTP). Unlike the previous work, our work can make the pixels in the edge position combine with the context information and carry out effective spatial displacement to achieve the natural effect of the alignment of the boundary position.

2.2 Multi-scale Features Fusion

Feature maps of different scales have different information. Shallow features have more details, while in-depth feature semantic information is richer. The advantage of multi-scale feature fusion is that features can fuse information of different scales, enrich semantic features and detail features, and effectively distinguish objects of different sizes in an image. In object detection, Feature Pyramid Networks(FPN) [38] firstly propose a multi-scale fusion strategy to fuse multi-scale features gradually. Semantic-FPN [29] and SETR-MLA [66] extend FPN and adopt top-down fusion mode to fuse multi-scale features. Based on this top-down fusion, ZigZagNet [35] propose top-down and bottom-up propagations to aggregate multi-scale features. Differently, DeepLab series [5, 8] introduces the dilated convolution, which fuses multi-scale features via concatenation at the channel dimension. Zhao et al. [65] introduces the Pyramid Pooling Module for context modeling. Zhu et al. [69] uses a non-local approach to fuse features from two different scales. Lin et al. [36, 37] uses a sparse sampling strategy to fuse multi-scale features. Zhang et al. [63] proposes a Feature Pyramid Transformer for multi-scale feature fusion, which adopts the feature map of all scales to the same scale. Grounding Transformer is used to carry out the fusion of features, abandoning the traditional top-down and bottom-up pathway. Different from the previous methods, we explore how to select the features of all scales at one time for global modeling based on retaining scale-level information (i.e., fusion without sampling to the same size), which is computationally friendly.

2.3 Transformer in Semantic Segmentation

Encouraged by Transformer’s great success in the field of natural language processing, Transformer has also made great strides in the field of image recognition. Vision Transformer(ViT) [16] is the first model to apply Transformer to image classification and achieve state-of-the-art. Swin Transformer [41] uses a pyramid structure, proposing a shifted window-based self-attention and performing well in both classified and downstream tasks. Recently, Transformer has received more and more attention on semantic segmentation [11, 66]. Since Transformer naturally has a global receptive field, this is friendly to segmentation tasks. Zheng [66] is the first semantic segmentation model using ViT as the backbone, which uses a Transformer-based encoder and CNN-based decoder. Ranftl et al. [45] also use ViT as the backbone and uses the convolutional progressive fusion feature to get full-resolution predictions. Many works are beginning to apply Transformer to decoder architecture. Xie et al. [57] uses top features extracted from the encoder. K types of learnable Token Embedded are randomly generated for transparent object segmentation. Nevertheless, Transformer is less suited to semantic segmentation when sequences are long. Therefore, SegFormer [58] uses the Transformer-based encoder to extract features and the lightweight MLP-decoder to predict pixel by pixel. FTN [53] and Pale Transformer [54] use grouping strategies to fuse context features. The Transformer is not friendly to sequence length, so Segmenter [48] uses features from the top extracted from the encoder and generates k classes of learnable tokens embedded in the Transformer for segmentation. Cheng et al. [11] build a unified framework for instance and semantic segmentation by introducing mask predictions. Although Transformer has the feature of global modeling, it is limited by the input of large resolution images(i.e., self-attention). So, we use a sparse sampling strategy to sample from multi-scale features with sampling points far smaller than feature-sequence, and then cross-attention is used to conduct cross-scale global modeling.

3.1 Framework

The overall framework of our CSSFormer is shown in Fig. 2, which consists of a Channel Information Augmentation module(CIA), Multi-Scale Feature Enhancement module (MSFE), and Cross-Scale Fusion module (CSF). Recently, Swin Transformer’s [41] powerful modeling capabilities in downstream tasks are impressive, so we chose it as the backbone of our method. Given an input image I ∈ ℝ 3 × H × W , We first use a backbone to get a multi-scale feature set, and use a top-down fusion approach similar to FPN [38] to deliver strong semantics to the bottom layer, which has little computational overhead, and resulting in pyramid features { F i ∈ ℝ C × H 2 i + 1 × W 2 i + 1 } i = 1 4 with richer semantics, where i, C, H, W denotes the stage index of the backbone, channel number, height and width, respectively. Then, we combine the obtained multi-scale features, use the Multi-Scale Feature Enhancement module to correct the deviation of the sampling points of multi-scale features, and carry out the fusion of multi-scale features to obtain a multi-scale feature sequence with more robust semantics and more accurate spatial position. The Channel Information Augmentation module is then used to enhance the features of each query feature, enrich the semantics of each query feature, and assign different weights to the channels adaptively. Next, query and multi-scale features are fed into the Cross-Scale Fusion module, and cross-attention is used for the final Cross-Scale Fusion. When each query feature has all-scale feature information, it also has a global receptive field. The spatial position of the sampling points of each query feature is corrected to obtain the final multi-scale feature set. Finally, a top-down pathway is adopted to fuse multi-scale feature sets and output them.

OPEN IN VIEWER

Fig. 2

Overall architecture of our Cross-Scale Sampling Transformer(CSSFormer). "CIA," "MSFE," and "CSF" indicate Channel Information Augmentation module, Multi-Scale Feature Enhancement module, and Cross-Scale Fusion module, respectively.

In order to facilitate the subsequent recalibration of query feature sampling points, details of query features are required to be highlighted. We hope to adaptively adjust the importance degree of each channel relative to the global context features so that the calibration can be carried out dynamically with the sampling points. We find that the richer the acceptance field of the feature map is, the more conducive it is to the expansion and contraction of channel weights and the stronger the responsiveness of sampling points. To this end, we design a Channel Information Augmentation module to re-model the query feature.

The data flow diagram of the entire Channel Information Augmentation module is shown in Fig. 3. Specifically, We use adaptive averaging pooling to obtain feature maps of four different sizes for each scale feature { F i ∈ ℝ C × H 2 i + 1 × W 2 i + 1 } i = 1 4 . This process is formulated as follow:

OPEN IN VIEWER

Fig. 3

Channel Information Augmentation module. "POOL" means global average pooling. "CONV" represents the 1×1 convolution layer. Φ means upsampling operation. φ means 3×3 convolution layer and global average pooling. "MLP" consists of several fully connected layers and a sigmoid function.

g i , j = { Average pooling ( F i ) } j = 1 4 ,(1)

In addition, 1×1 convolution is used to align the number of channels and features with different receptive fields are obtained, and upsampling is performed to the size of the query feature for fusion.

g i , j ′ = Conv ( g i , j ) ,(2)

F i ′ = concatenate ( F i , Φ ( g i , j ′ ) ) ,(3)

For significance feature F i ′ , we obtain a global feature vector through 3×3 convolution layer and global average pooling. This process is formulated as follows:

α i = φ ( F i ′ ) .(4)

Next, for global feature vector α _i , we use an MLP module (i.e. two 1 × 1 conv layers followed by a sigmoid activation function, the input and output dimensions of the two 1 × 1 conv layers change as C->C/2 and C/2 ->C) to predict the importance of each channel. This process is formulated as follows:

α i ' = MLP ( a ) .(5)

To get the importance of each channel on the feature map, we also need to dot and F i to get the final output. This process is formulated as follows:

F i ′ ˆ = α i ′ × F i + F i .(6)

It is worth mentioning that our Channel Information Augmentation module was inspired by SENet [24] and PSPNet [65]. The difference is that we hope to more effectively scale the channel weights by enhancing the receptive field of the feature map to correspond to the sampling point and feel the farther area and residual connections are used to prevent the channel information from being overly squeezed and excited. In this way, local details will be highlighted more after the CIA Module is used. The query feature will be re-modeled with only a little computation, conducive to the subsequent multi-scale information fusion and adjustment of edge detail features.

We find that simply superimposing multi-scale information is not friendly to the query feature of each scale, and it cannot effectively obtain information gained from it. Therefore, the MSFE module (as shown in Fig. 4) is proposed to enhance multi-scale information, mainly through the interaction of multi-scale features, fine-tuning the spatial position of the pixel in the spatial domain, and one-time fusion of multi-scale information. We find that the one-time integration of all information scales would produce a considerable amount of computation and cause a significant GPU memory overhead, resulting in difficult training. Inspired by Deformable detr [68], we adopt a sparse sampling strategy to solve this problem(as shown in Fig. 5).

OPEN IN VIEWER

Fig. 4

Multi-Scale Feature Enhancement module. Sampling offset requires only one fully connected layer to generate.

OPEN IN VIEWER

Fig. 5

Pyramid Feature Sampling Alignment(PFSA). The pyramid features are sampled by sampling points and aligned by sampling offset.

Specifically, we flatten the multi-scale features { F i ∈ ℝ C × H 2 i + 1 × W 2 i + 1 } i = 1 4 and simply stack them together to obtain a full-scale feature sequence, denoted as F ˆ ∈ ℝ C × ∑ i = 1 4 H 2 i + 1 × W 2 i + 1 . Similar to the traditional self-attention mechanism, we also map query and value through two fully connected layers, except that keys are discarded. The Attention weights are obtained directly from query, reducing Transformer’s second complexity to one. Specifically:

A = Softmax ( ∑ i = 1 k Weight i ) × Value ( C s + ▵ i ) .(7) Where k is the number of sparse sampling pixels, Weight _i is the attention weights for the i-th sampled key element. C _s is the coordinates of sampling pixels and ▵ _i denote the interpolation method and sampling offsets. The generation of sampling offsets can be obtained only through a fully connected layer.

The importance of effective modeling of multi-scale information has been demonstrated in [8, 69]. As shown in Fig. 1, we hope to avoid the position offset caused by upsampling in the top-down fusion process [29, 45]. And to be able to merge the full-scale features obtained from the MSFE module at one time. As shown in Fig. 7, we propose our CSF module, which can calibrate the spatial position of pixels and obtain the full-scale information under the action of multi-scale information. For each scale feature, we carry out Multi-Head Self-Attention on F i , which highlights the feature with strong semantics and is conducive to the generation of sampling offset in the future. In addition, we set up a learnable matrix α ∈ ℝ 1 × C , which can dynamically scale and shrink the Attention weight according to the network learning to avoid too large or too small weight. In particular:

OPEN IN VIEWER

Fig. 6

Cross-Scale Fusion module."RP" means reference point. Value is the output from the MSFE Module.

OPEN IN VIEWER

Fig. 7

Visual analysis.

A ˆ = α × Softmax ( qk T d head ) v ,(8) In this way, the query feature is obtained, and the value feature is obtained from the MSFE module. Let us rewrite formula (7):

Attention = Softmax ( ∑ i = 1 k Weight qi ) × A ( C s + ▵ i ) .(9) It is worth mentioning that the Weight _qi is from the query feature, and the generation of ▵ _i is from the query feature in order to feature-aligned multi-scale information through the query feature. A is the output from formula (7).

A multi-task loss function of AUX and CE is used to optimize the model parameters jointly. In particular, the loss function of AUX is defined as:

L AUX = 1 H × W ∑ L ce ( F 3 K × H × W , G T ) .(10)

Here, we use Swin Transformer as the backbone and use L AUX for the third stage feature, L _ce and G T represent the cross entropy loss and ground truth, respectively. F 3 represents 16 times of upsampling on the feature map of the encoder’s third stage. For the final per-pixel classification, we define the loss function of L CE as:

L CE = 1 H × W ∑ L ce ( O , G T ) .(11)

Here, O indicates that the output feature is upsampled eight times to obtain the size of the original figure.

Finally, we formulate the multi-task loss function L as:

L = λ L AUX + L CE .(12) We empirically set λ = 0.4 by default. The model parameters are learned jointly through back propagation with this joint loss function.

4.1 Datasets

PASCAL Context [44] is an extension of the PASCAL VOC 2010 detection challenge. It contains 4998 and 5105 images for training and validation, respectively. We evaluated the most frequent 60 classes (59 categories with background) following previous works.

ADE20K [67] is a challenging benchmark, including 150 categories and diverse scenes with 1,038 image-level labels, split into 20000 and 2000 images for training and validation.

COCO-Stuff 10K [3] is a significant scene parsing benchmark with 9000 training images and 1000 testing images with 182 categories (80 objects and 91 stuff).

Cityscapes [13] carefully annotates 19 object categories of urban landscape images. It contains 5K finely annotated images, split into 2975 and 500 for training and validation.

4.2 Implementation details

Because Vision Transformer has been prominent in the visual field recently, and the community is gradually using Transformer to replace CNN as the backbone. Since Swin Transformer [41] is a pyramid structure with low computation overhead, we choose it as our backbone. The channel C of features F i is set to 256. We select k=16 sampling points for the feature map of each scale. Both the MSFE module and CSF module have a depth of 2. The head number of the MSFE module and CSF module is set to 12. Following the default setting (e.g., data augmentation and training schedule) of public codebase mmsegmentation [12]. We optimize our models using AdamW [43] with a batch size of 16 (batch size of 8 for Cityscapes), and we adopt a polynomial learning rate decay schedule. Momentum and weight decay are set to 0.9 and 0.01, respectively, for all the experiments on the four datasets. We set the initial learning rate at 0.00006 on ADE20K and Cityscapes and 0.00002 on Pascal Context and COCO-Stuff 10K. The total iterations are set to 160k, 60k, 80k, and 80k for ADE20K, COCO-Stuff 10K, Cityscapes, and PASCAL Context, respectively. During training, data augmentation in all the experiments consists of three steps:(i) random horizontal flipping, (ii) we apply random resize with the ratio between 0.5 and 2, (iii) random cropping (512x512, 512x512, 768x768, and 480x480 for ADE20K, COCO-Stuff 10K, Cityscapes and Pascal Context respectively). For a fair comparison, we apply the synchronized BN to the decoder and auxiliary loss head(encoder-stage3), and we do not adopt the widely-used tricks such as OHEM [47] loss, different learning rate(LR) for backbone and decoder head, and class balanced loss in model training. During inference, we use the default settings of mmsegmentation [12]. Precisely, the input image is first scaled to a uniform size. Multi-scale scaling and random horizontal flip are then performed on the image with scaling factors (0.5, 0.75, 1.0, 1.25, 1.5, 1.75). The sliding window is adopted for the test (e.g., 512 × 512 for ADE20K).

4.3 Comparisons with the State-of-the-art Methods

Results on ADE20K. Table 1 shows the results on ADE20K dataset. We compare the performance of CSSFormer with state-of-the-art methods. For previous work, ISNet [28] achieves mIoU 47.55% by using CNN as the backbone. SegFormer [58] and MCIBI [27] achieves mIoU 51.80%, 50.80% respectively by using Transformer as backbone. Our CSSFormer achieves mIoU of 53.12% with single-scale (SS) inference. When multi-scale inference is adopted, our method achieves a new state-of-the-art with mIoU hitting 54.02% and +2.22 mIoU higher than SegFormer (54.02 vs. 51.80). When training with larger image sizes (640 × 640), our method is +1.88 mIoU higher than Segmenter [48] (55.48 vs. 53.60). To be fair, we did not use any tricks in our experiment.

Results on Pascal Context. Table 2 compares the segmentation results on Pascal Context. GINet [55] with ResNet-101 backbone achieves a mIoU of 54.90%, and SR [62] with HRNet-w48 backbone achieves 55.70 mIoU. When using Transformer as the backbone, the previous best method, Upernet achieves mIoU 57.29%. Our CSSFormer achieves a superior mIoU of 59.95%, which is +2.66 mIoU higher than UperNet (59.95 vs. 57.29).

Results on COCO-Stuff 10K. The state-of-the-art results on COCO-Stuff 10K dataset are shown in Table 3. ISNet with ResNeSt-101 as the backbone achieves mIoU 42.08%. MCIBI and Upernet achieve mIoU 44.89%, 47.71% respectively. Our CSSFormer achieves a superior mIoU of 50.92% with multi-scale inference.

Results on Cityscapes. Tables 4 show the comparative results on the Cityscapes validation set. We can see that our CSSFormer achieves a new state of the art with mIoU(SS) hitting 83.78%, which is +4.44% higher than SETR [66] (83.78 vs. 79.34), for multi-scale testing, our method is +2.57 mIoU(MS) higher than SETR(84.72 vs. 82.15). For Segformer trained with higher resolution(1024×1024 is much larger than our crop size 768×768), CSSFormer is still +1.38% mIoU(SS) higher than it (83.78 vs. 82.40), for multi-scale testing, our method is +0.72 mIoU(MS) higher than SegFormer(84.72 vs. 84.00).