S EA F ORMER : S QUEEZE - ENHANCED A XIAL T RANSFORMER FOR M OBILE S EMANTIC S EGMENTATION
A BSTRACT
Since the introduction of Vision Transformers, the landscape of many computer vision tasks ( e.g. , semantic segmentation), which has been overwhelmingly dominated by CNNs, recently has significantly revolutionized.
However, the computational cost and memory requirement render these methods unsuitable on the mobile device, especially for the high-resolution per-pixel semantic segmentation task.
In this paper, we introduce a new method squeeze-enhanced Axial Transformer ( SeaFormer ) for mobile semantic segmentation.
Specifically, we design a generic attention block characterized by the formulation of squeeze Axial and detail enhancement.
It can be further used to create a family of backbone architectures with superior cost-effectiveness.
Coupled with a light segmentation head, we achieve the best trade-off between segmentation accuracy and latency on the ARM-based mobile devices on the ADE20K and Cityscapes datasets.
Critically, we beat both the mobile-friendly rivals and Transformer-based counterparts with better performance and lower latency without bells and whistles.
Beyond semantic segmentation, we further apply the proposed SeaFormer architecture to image classification problem, demonstrating the potentials of serving as a versatile mobile-friendly backbone.
Our code and models are made publicly available at https://github.com/fudan-zvg/SeaFormer .
I NTRODUCTION
As a fundamental problem in computer vision, semantic segmentation aims to assign a semantic class label to each pixel in an image.
Conventional methods rely on stacking local convolution kernel Long et al.
(2015) to perceive the long-range structure information of the image.
Since the introduction of Vision Transformers Dosovitskiy et al.
(2021), the landscape of semantic segmentation has significantly revolutionized.
Transformer-based approaches Zheng et al.
Xie et al.
(2021) have remarkably demonstrated the capability of global context modeling.
However, the computational cost and memory requirement of Transformer render these methods unsuitable on mobile devices, especially for high-resolution imagery inputs.
Following conventional wisdom of efficient operation, local/window-based attention Luong et al.
; Liu et al.
; Huang et al.
(2021a)
; Yuan et al.
(2021), Axial attention Huang et al.
(2019b)
; Ho et al.
; Wang et al.
(2020a), dynamic graph message passing Zhang et al.
(2020; 2022b) and some lightweight attention mechanisms Hou et al. ; Li et al. (2021b;c; 2020); Liu et al.
; Shen et al.
; Xu et al.
; Cao et al.
; Woo et al.
Wang et al.
(2020b); Choromanski et al.
; Chen et al.
(2017); Mehta & Rastegari (2022a) are introduced.
However, these advances are still insufficient to satisfy the design requirements and constraints for mobile devices due to the high latency on the high-resolution inputs (see Figure 1).
Recently there is a surge of interest in building a Transformer-based semantic segmentation.
In order to reduce the computation cost at high resolution, TopFormer Zhang et al.
(2022c) dedicates to applying the global attention at a 1 / 64 scale of the original input, which definitely harms the segmentation performance.
To solve the dilemma of high-resolution computation for pixel-wise segmentation task and low latency requirement on the mobile device in a performance harmless way, we propose a family
[Figure]
of mobile-friendly Transformer-based semantic segmentation model, dubbed as squeeze-enhanced
Axial Transformer (SeaFormer), which reduces the computational complexity of axial attention from
O (( H + W ) HW ) to O ( HW ) , to achieve superior accuracy-efficiency trade-off on mobile devices and fill the vacancy of mobile-friendly efficient Transformer.
The core building block squeeze-enhanced Axial attention (SEA attention) seeks to squeeze (pool) the input feature maps along the horizontal/vertical axis into a compact column/row and computes self-attention.
We concatenate query, keys and values to compensate the detail information sacrificed during squeeze and then feed it into a depth-wise convolution layer to enhance local details.
Coupled with a light segmentation head, our design (see Figure 2) with the proposed SeaFormer layer in the small-scale feature is capable of conducting high-resolution image semantic segmentation with low latency on the mobile device.
As shown in Figure 1, the proposed SeaFormer outperforms other efficient neural networks on the ADE20K dataset with lower latency.
In particular,
SeaFormer-Base is superior to the lightweight CNN counterpart MobileNetV3 (41.0 vs. 33.1 mIoU) with lower latency (106ms vs. 126ms) on an ARM-based mobile device.
We make the following contributions : (i)
We introduce a novel squeeze-enhanced Axial Transformer (SeaFormer) framework for mobile semantic segmentation; (ii) Critically, we design a generic attention block characterized by the formulation of squeeze Axial and detail enhancement;
It can be used to create a family of backbone architectures with superior cost-effectiveness; (iii) We show top performance on the ADE20K and Cityscapes datasets, beating both the mobile-friendly rival and Transformer-based segmentation model with clear margins; (iv) Beyond semantic segmentation, we further apply the proposed SeaFormer architecture to the image classification problem, demonstrating the potential of serving as a versatile mobile-friendly backbone.
R ELATED WORK
Combination of Transformers and convolution Convolution is relatively efficient but not suitable to capture long-range dependencies and vision Transformer has the powerful capability with a global receptive field but lacks efficiency due to the computation of self-attention.
In order to make full use of both of their advantages, MobileViT Mehta & Rastegari (2022a), TopFormer Zhang et al.
(2022c), LVT Yang et al.
Mobile-Former Chen et al.
(2022b), EdgeViTs Pan et al.
(2022a),
MobileViTv2 Mehta & Rastegari (2022b), EdgeFormer Zhang et al.
(2022a)
and
EfficientFormer Li et al.
(2022) are constructed as efficient ViTs by combining convolution with Transformers.
Mobile-
ViT, Mobile-Former, TopFormer and EfficientFormer are restricted by Transformer blocks and have to trade off between efficiency and performance in model design.
LVT, MobileViTv2 and EdgeViTs keep the model size small at the cost of relatively high computation, which also means high latency.
Axial attention and variants Axial attention Huang et al.
(2019b)
; Ho et al.
Wang et al.
(2020a) is designed to reduce the computational complexity of original global self-attention Vaswani et al.
It computes self-attention over a single axis at a time and stacks a horizontal and a vertical axial attention module to obtain the global receptive field.
Strip pooling Hou et al.
and Coordinate attention Hou et al.
(2021) uses a band shape pooling window to pool along either the horizontal or the vertical dimension to gather long-range context.
Kronecker Attention Networks Gao et al.
(2020) uses the juxtaposition of horizontal and lateral average matrices to average the input matrices and performs attention operation.
These methods and other similar implementations provide performance gains partly at considerably low computational cost compared with Axial attention.
However, they ignore the lack of local details brought by the pooling/average operation.
Mobile semantic segmentation The mainstream of efficient segmentation methods are based on lightweight CNNs.
DFANet Li et al.
(2019) adopts a lightweight backbone to reduce computation cost and adds a feature aggregation module to refine high-level and low-level features.
ICNet Zhao et al.
(2018) designs an image cascade network to speed up the algorithm, while BiSeNet Yu et al.
(2018; 2021) proposes two-stream paths for low-level details and high-level context information, separately.
Fast-SCNN Poudel
et al.
(2019) shares the computational cost of the multi-branch network to yield a run-time fast segmentation CNN.
TopFormer Zhang et al.
(2022c) presents a new architecture with a combination of CNNs and ViT and achieves a good trade-off between accuracy and computational cost for mobile semantic segmentation.
However, it is still restricted by the heavy computation load of global self-attention.
M ETHOD 3. O VERALL ARCHITECTURE
Inspired by the two-branch architectures Yu et al.
; Poudel et al.
; Hong et al.
Huang et al.
(2021b)
; Chen et al.
(2022b)
, we design a s queeze-
e nhanced A xial Trans former ( SeaFormer ) framework.
As is shown in Figure 2, SeaFormer consists of these parts: shared STEM , context branch , spatial branch , fusion block and light segmentation head .
For a fair comparison, we follow TopFormer Zhang et al.
(2022c) to design the STEM.
It consists of one regular convolution with stride of 2 followed by four MobileNet blocks where stride of the first and third block is 2.
The context branch and the spatial branch share the produced feature map, which allows us to build a fast semantic segmentation model.
Context branch The context branch is designed to capture context-rich information from the feature map x s .
As illustrated in the red branch of Figure 2, the context branch is divided into three stages.
To obtain larger receptive field, we stack SeaFormer layers after applying a MobileNet block to down-sampling and expanding feature dimension.
Compared with the standard convolution as the down-sampling module, MobileNet block increases the representation capacity of the model while maintaining a lower amount of computation and latency.
For variants except SeaFormer-Large,
SeaFormer layers are applied in the last two stages for superior trade-off between accuracy and efficiency.
For SeaFormer-Large, we insert SeaFormer layers in each stage of context branch.
To achieve a good trade-off between segmentation accuracy and inference speed, we design a squeezeenhanced Axial attention block (SEA attention) illustrated in the next subsection.
Spatial branch The spatial branch is designed to obtain spatial information in high resolution.
Identical to the context branch, the spatial branch reuses feature maps x s .
However, the feature from the early convolution layers contains rich spatial details but lacks high-level semantic information.
Consequently, we design a fusion block to fuse the features in the context branch into the spatial branch, bringing high-level semantic information into the low-level spatial information.
Fusion block As depicted in Figure 2, high resolution feature maps in the spatial branch are followed by a 1 × 1 convolution and a batch normalization layer to produce a feature to fuse.
Low
[Figure]
resolution feature maps in the context branch are fed into a 1 × 1 convolution layer, a batch normalization layer, a sigmoid layer and up-sampled to high resolution to produce semantics weights by bilinear interpolation.
Then, the semantics weights from context branch are element-wisely multiplied to the high resolution feature from spatial branch.
The fusion block enables low-level spatial features to obtain high-level semantic information.
segmentation head The feature after the last fusion block is fed into the proposed segmentation head directly, as demonstrated in Figure 2.
For fast inference purpose, our light segmentation head consists of two convolution layers, which are followed by a batch normalization layer separately and the feature from the first batch normalization layer is fed into an activation layer.
3. S QUEEZE - ENHANCED A XIAL ATTENTION
The global attention can be expressed as
[Equation]
1) where x ∈ R H × W × C .
q , k , v are linear projection of x , i.e. q = W q x , k = W k x , v = W v x , where W q , W k ∈ R C qk × C , W v ∈ R C v × C are learnable weights.
G ( o ) means all positions on the feature map of location o
= ( i, j ) .
When traditional attention module is applied on a feature map of
H × W × C , the time complexity can be O
( H 2 W 2 ( C qk + C v )) , leading to low efficiency and high latency.
[Equation]
To improve the efficiency, there are some works Liu et al.
; Huang et al.
(2019b)
; Ho et al.
(2019) computing self-attention within the local region.
We show two most representative efficient
Transformer in Equation 2, 3.
Equation 2 is represented by window-based attention Luong et al.
(2015) successfully reducing the time complexity to O ( m 2 HW ( C qk + C v ))
=
O ( HW ) , where
N m × m ( o ) means the neighbour m × m positions of o , but loosing global receptiveness.
The Equation 3 is represented by Axial attention Ho et al.
(2019), which only reduces the time complexity to O
(( H + W ) HW ( C qk + C v ))
=
O (( HW ) 1 . 5 ) , where N H × 1 ( o ) means all the positions of the column of o ; N 1 × W ( o ) means all the positions of the row of o .
According to their drawbacks, we propose the mobile-friendly squeeze-enhanced Axial attention, with a succinct squeeze Axial attention for global semantics extraction and an efficient convolutionbased detail enhancement kernel for local details supplement. [Equation]
[Equation]
𝐾𝑉 𝑄
Concat
Multi-head attention
Squeeze Axial attention
Broadcast
Broadcast 3 x 3 d w c o n v
B N
B N R e L U 6 1 x 1 C o n v
Multi-headattention 1 x 1 C o n v
Horizontalsqueeze
Verticalsqueeze
Detail enhancement kernel 𝑄
ℎ 𝐾
ℎ 𝑉 ℎ 𝑄
𝑣 𝐾 𝑣 𝑉 𝑣
Squeeze-enhancedAxial attention
FFN 𝑯×𝑾×𝑪
𝒒𝒌 𝑥
Mul 𝑯×𝑾×(𝟐𝑪 𝒒𝒌 +𝑪 𝒗 ) 𝑯×𝟏×𝑪 𝒒𝒌 𝑯×𝑾×𝑪 𝒗 𝑯×𝟏×𝑪 𝒗 𝟏×𝑾×𝑪 𝒒𝒌 𝟏×𝑾×𝑪 𝒗
Squeeze Axial attention To achieve a more efficient computation and aggregate global information at the same time, we resort to a more radical strategy.
In the same way, q , k , v are first get from x with W ( s ) q , W ( s ) k ∈ R C qk × C , W ( s ) v ∈ R C v × C .
According to Equation 4, we first implement horizontal squeeze by taking average of query feature map on the horizontal direction.
In the same way, the right shows the vertical squeeze on the vertical direction.
z → ( · ) means permuting the dimension of tensor z as given, and 1 m ∈ R m is a vector with all the elements equal to 1.
The squeeze operation on q also repeats on k and v , so we finally get q ( h ) , k ( h ) , v ( h ) ∈
R H
× C qk , q ( v ) , k ( v ) , v ( v ) ∈ R W × C
qk .
The squeeze operation reserves the global information to a single axis, thus greatly alleviating the following global semantic extraction showing by Equation 5.
[Equation]
Each position of feature map propagates information only on two squeezed axial features.
Although it shows no distinct computation reduction comparing to Equation 3, repeat of Equation 5 can be simply implemented by the most efficient broadcast operation.
The detail is shown in Figure 3.
Time complexity for squeezing q , k , v is O
(( H + W )(2 C qk + C v ))
and the attention operation takes O
(( H 2 + W 2 )( C qk + C v )) time.
Thus, our squeeze Axial attention successfully reduces time complexity to O ( HW ) .
Squeeze Axial position embedding Equation 4 are, however, not positional-aware, including no positional information of feature map.
Hence, we propose squeeze Axial position embedding to squeeze Axial attention.
For squeeze Axial attention, we render both q ( h ) and k ( h ) to be aware of their position in squeezed axial feature by introducing positional embedding r q ( h ) ,
r k ( h ) ∈
R H
× C qk , which are linearly interpolated from learnable parameters B q ( h ) ,
B k ( h ) ∈ R L × C
qk .
L is a constant.
In the same way, r q ( v ) , r k ( v ) ∈ R W × C qk are applied to q ( v ) , k ( v ) .
Thus, the positional-aware squeeze
Axial attention can be expressed as Equation 6.
[Equation]
[Equation]
Detail enhancement kernel The squeeze operation, though extracting global semantic information efficiently, sacrifices the local details.
Hence an auxiliary convolution-based kernel is applied to enhance the spatial details.
As is shown in the upper path of Figure 3, q , k , v are first get from x with another W ( e ) q , W ( e ) k ∈ R C qk × C , W ( e ) v ∈ R C v × C and are concatenated on the channel dimension and then passed to a block made up of 3 × 3 depth-wise convolution and batch normalization.
[Table]
By using a 3 × 3 convolution
[Table]
auxiliary local details can be aggregated from q
[Table]
k
[Table]
v .
And then a linear projection with activation function and batch normalization are used to squeeze (2 C qk + C v ) dimension to C and generate detail enhancement weights.
Finally, the enhancement feature will be fused with the feature given by squeeze Axial attention.
Different enhancement mode including element-wise addition and multiplication will be compared in experiment section.
Time complexity for the 3 × 3 depth-wise convolution is O (3 2 HW (2 C qk + C v )) and the time complexity for the 1 ×convolution is O ( HWC (2 C qk + C v )) .
Time for the other operations like activation can be omitted.
Architecture and Variants We introduce four variants, SeaFormer-Tiny, Small, Base and Large (T, S, B and L).
More configuration details are listed in the supplementary material.
E XPERIMENTS
We evaluate our method on semantic segmentation and image classification tasks.
First, we describe implementation details and compare results with state of the art.
We then conduct a series of ablation studies to validate the design of SeaFormer.
Each proposed component and important hyper-parameters are examined thoroughly.
4. E XPERIMENTAL SETUP 4.1. D ATASET
We perform segmentation experiments over ADE20K Zhou et al.
, CityScapes Cordts et al.
The mean of intersection over union (mIoU) is set as the evaluation metric.
We convert full-precision models to TNN Contributors (2019) and measure latency on an ARM-based device with a single Qualcomm Snapdragon 865 processor.
[Table] ADE20K dataset covers 150 categories, containing 25K images that are split into 20K/2K/3K for Train , val and test .
CityScapes is a driving dataset for semantic segmentation.
It consists offine annotated high-resolution images with 19 categories.
4.1. I MPLEMENTATION DETAILS
We set ImageNet-1K Deng et al.
(2009) pretrained network as the backbone, and training details of
ImageNet-1K are illustrated in the last subsection.
For semantic segmentation, the standard BatchNorm Ioffe & Szegedy (2015) layer is replaced by synchronized BatchNorm.
Training Our implementation is based on public codebase mmsegmentation Contributors (2020).
We follow the batch size, training iteration scheduler and data augmentation strategy of
TopFormer Zhang et al.
(2022c) for a fair comparison.
The initial learning rate is 0.0005 and the weight decay is 0.01.
A “poly” learning rate scheduled with factor 1.0 is adopted.
During inference, we set the same resize and crop rules as TopFormer to ensure fairness.
The comparison of Cityscapes contains full-resolution and half-resolution.
For the full-resolution version, the training images are randomly scaled and then cropped to the fixed size of 1024 × 1024.
For the half-resolution version, the training images are resized to 1024 × 512 and randomly scaling, the crop size is 1024 × 512.
4. C OMPARISON WITH STATE OF THE ART
ADE20K Table 1 shows the results of SeaFormer and previous efficient backbones on ADE20K val set.
The comparison covers Params, FLOPs, Latency and mIoU. As shown in Table 1, SeaFormer outperforms these approaches with comparable or less FLOPs and lower latency.
Compared with specially designed mobile backbone, TopFormer, which sets global attention as its semantics extractor, SeaFormer achieves higher segmentation accuracy with lower latency.
And the performance of SeaFormer-B surpasses MobileNetV3 by a large margin of +7.9% mIoU with lower latency (-16%).
The results demonstrate our SeaFormer layers improve the representation capability significantly.
Cityscapes From the table 2, it can be seen that SeaFormer-S achieves comparable or better results than TopFormer-B with less computation cost and latency, which proves that SeaFormer could also achieve a good trade-off between performance and latency in high-resolution scenario.
4. A BLATION STUDIES
In this section, we ablate different self-attention implementations and some important design elements in the proposed model, including our squeeze-enhanced Axial attention module (SEA attention) and fusion block on ADE20K dataset.
The influence of components in SEA attention We conduct experiments with several configurations, including detail enhancement kernel only, squeeze Axial attention only, and the fusion of both.
As is shown in Table 3, only detail enhancement or squeeze Axial attention achieves a relatively poor
[Table]
performance, and enhancing squeeze Axial attention with detail enhancement kernel brings a per formance boost with a gain of 2.3% mIoU on ADE20K.
The results that enhancing global semantic features from squeeze Axial attention with local details from convolution optimizes the feature extraction capability of Transformer block.
For enhancement input, there is an apparent performance gap between upconv(x) and conv(x).
And we conclude that increasing the channels will boost performance significantly.
Comparing concat[qkv] and upconv(x), which also correspond to w/ or w/o convolution weight sharing between detail enhancement kernel and squeeze Axial attention, we can find that sharing weights makes our model improve inference efficiency with minimal performance loss (35.8 vs. 35.9).
As for enhancement modes, multiplying features from squeeze
Axial attention and detail enhancement kernel outperforms add enhancement by +0.4% mIoU.
[Table]
Comparison with different selfattention modules In order to eliminate the impact of our architecture and demonstrate the effectiveness and generalization ability of SEA attention, we ran experiments on Swin Transformer Liu et al.
(2021) by replacing window attention in Swin Transformer with different attention blocks.
We set the same training protocol, hyperparameters, and model architecture configurations as Swin for a fair comparison.
When replacing window attention with CCAttention (CCNet) or DoubleAttention (A2-Nets), they have much lower FLOPs than SeaFormer and other attention blocks.
Considering that we may not be able to draw conclusions rigorously, we doubled the number of their Transformer blocks (including MLP).
As ACmix has the same architecture configuration as Swin, we borrow the results from the original paper.
From Table 4, it can be seen that SeaFormer outperforms other attention mechanisms with lower FLOPs and latency.
The influence of the width in fusion block To study the influence of the width in fusion block, we perform experiments with different embedding dimensions in fusion blocks on SeaFormer-Base, M denotes the channels that spatial branch and context branch features mapping to in two fusion blocks.
Results are shown in Table 5.
4. I MAGE CLASSIFICATION
We conduct experiments on ImageNet-1K Deng et al.
(2009), which contains 1.28M training images and 50K validation images from 1,000 classes.
We employ an AdamW Kingma & Ba (2014) optimizer for 600 epochs using a cosine decay learning rate scheduler.
A batch size of 1024, an initial learning rate of 0.064, and a weight decay of 2e-5 are used.
The results are illustrated in Table 6.
Compared with other efficient approaches, SeaFormer achieves a relatively better trade-off between latency and accuracy. [Table]
4. L ATENCY STATISTICS
[Table] We make the statistics of the latency of the proposed SeaFormer-Tiny, as shown in Figure 4, the shared STEM takes up about half of the latency of the whole network (49%).
The latency of the context branch is about a third of the total latency (34%), whilst the actual latency of the spatial branch is relatively low (8%) due to sharing early convolution layers with the context branch.
Our light segmentation head (8%) also contributes to the success of building a light model.
C ONCLUSION
In this paper, we have proposed squeeze-enhanced Axial Transformer ( SeaFormer ) for mobile semantic segmentation, filling the vacancy of mobile-friendly efficient Transformer.
Moreover, we create a family of backbone architectures of SeaFormer and achieve cost-effectiveness.
The superior performance on the ADE20K and Cityscapes, and the lowest latency demonstrate its effectiveness on the ARM-based mobile device.
Beyond semantic segmentation, we further apply the proposed SeaFormer architecture to image classification problem, demonstrating the potential of serving as a versatile mobile-friendly backbone. A CKNOWLEDGMENTS
This work was supported in part by National Natural Science Foundation of China (Grant No.
, Lingang Laboratory (Grant No. LG-QS-202202-07), Natural Science Foundation of
Shanghai (Grant No. 22ZR1407500) and CCF-Tencent Rhino-Bird Young Faculty Open Research
Fund (No. CCF-Tencent RAGR20210111).
Appendix A A RCHITECTURE DETAILS AND VARIANTS
SeaFormer backbone contains 6 stages, corresponding to the shared STEM and context branch in
Figure 2 in the main paper.
When conducting the image classification experiments, a pooling layer and a linear layer are added at the end of the context branch.
Table 7 details the family of our SeaFormer configurations with varying capacities.
We construct
SeaFormer-Tiny, SeaFormer-Small, SeaFormer-Base and SeaFormer-Large models with different scales via varying the number of SeaFormer layers and the feature dimensions.
We use input image size of 512 × 512 by default.
For variants except SeaFormer-Large, SeaFormer layers are applied in the last two stages for superior trade-off between accuracy and efficiency.
For SeaFormer-Large, we apply the proposed SeaFormer layers in each stage of the context branch.
B C OMPLEXITY ANALYSIS
we analyze the complexity of our proposed SEA attention in subsection 3.2 to demonstrate its efficiency theoretically.
In our application, we set C qk = 0 .
5 C v to further reduce computation cost.
The total time complexity of squeeze-enhanced Axial attention is
[Equation]
if we assume H = W and take channel as constant.
SEA attention is linear to the feature map size theoretically.
Moreover, SEA attention only includes mobile-friendly operation like convolution, pooling, matrix multiplication and so on.
C P ASCAL C ONTEXT PERFORMANCE
We evaluate performance on Pascal Context val set over 59 categories and 60 categories.
PASCAL Context dataset has 4998/5105 images for train and test , covering 59 semantic labels andbackground.
Following TopFormer Zhang et al.
(2022c)
, we train the models for 80,000 iterations on PASCAL
Context dataset.
The same data augmentation strategy and batch size are adopted for a fair comparison.
The initial learning rate is 0.0002 and the weight decay is 0.01.
A poly learning rare scheduled with factor 1.0 is used.
Table 8 demonstrates that SeaFormer-S is +1.4% mIoU higher (45.08% vs.43.68%) than TopFormerS with lower latency.
D COCO-S TUFF PERFORMANCE
We compare SeaFormer with the previous approaches on COCO-Stuff val set.
COCO-Stuff dataset augments COCO dataset with pixel-level stuff annotations.
10K complex images are selected from
COCO.
The train and test set contain 9K/1K images.
Following TopFormer Zhang et al.
(2022c)
, we train the models for 80,000 iterations on COCOStuff dataset.
The same data augmentation strategy and batch size are adopted for a fair comparison.
The initial learning rate is 0.0002 and the weight decay is 0.01.
A poly learning rare scheduled with factor 1.0 is used.
Table 9 reveals that SeaFormer-S is +2.0% mIoU higher (32.82% vs.30.83%) than TopFormer-S with less computation cost and lower latency. [Table] [Table] [Table]
[Table]
[Table]
E O BJECT DETECTION PERFORMANCE
To further demonstrate the generalization ability of our proposed SeaFormer backbone on other downstream tasks, we conduct object detection task on COCO dataset.
Setup We use RetinaNet Lin et al.
(2017) (one-stage) as the detection framework and follow the standard settings to use SeaFormer as backbone to generate e feature pyramid at multiple scales.
All models are trained on train2017 split for 12 epochs (1×) from ImageNet pretrained weights.
Results From the table 10 We can observe that our SeaFormer achieves superior results on detection task which further demonstrates the strong generalization ability of our method.
F A DDITIONAL ABLATION STUDY
In addition to the ablation study in the submission paper, we investigate the effect of fusion method in fusion block in Figure 2.
F. T HE INFLUENCE OF FUSION BLOCK DESIGN
We set four fusion methods, including ”Add directly”, ”Multiply directly”, ”Sigmoid add” and ”Sigmoid multiply”.
X directly means features from context branch and spatial branch X directly.
Sigmoid X means feature from context branch goes through a sigmoid layer and X feature from spatial branch.
From the Table 11 we can see that replacing sigmoid multiply with other fusion methods hurts performance.
Sigmoid multiply is our optimal fusion block choice.
F. E FFECTIVE AND EFFICIENCY OF SEA ATTENTION
To verify the effectiveness and efficiency of SEA attention based on our designed pipeline, we experiment with convolution, Global attention, Local attention, Axial attention and three convolution [Table] [Table] enhanced attention methods including our SEA attention, ACmix and MixFormer.
The ablation experiments are organized in seven groups.
Since the resolution of computing attention is relatively small, the window size in Local attention, ACmix, and MixFormer is set to 4.
We adjust the channels when applying different attention modules to keep the FLOPs aligned and compare their performance and latency.
The results are illustrated in Table 12.
As demonstrated in the table [Table] SEA attention outperforms the counterpart built on other efficient attentions.
Compared with global attention, SEA attention outperforms it by +1.2% Top1 accuracy on ImageNet-1K and +1.6 mIoU on ADE20K with less FLOPs and lower latency.
Compared with similar convolution enhanced attention works, ACmix and MixFormer, our SEA attention obtains better results on ImageNet-1K and ADE20K with similar FLOPs but lower latency.
The results indicate the effectiveness and efficiency of SEA attention module.
G P ERFORMANCE UNDER DIFFERENT PRECISION OF THE MODELS
Following TopFormer, we measure the latency in the submission paper on a single Qualcomm Snapdragon 865, and only an ARM CPU core is used for speed testing.
No other means of acceleration, e.g., GPU or quantification, is used.
We provide a more comprehensive comparison to demonstrate the necessity of our proposed method.
We test the latency under different precision of the models.
From the table 13, it can be seen that whether it is full precision or half precision the performance of SeaFormer is better than that of TopFormer.
H V ISUALIZATION H. A TTENTION HEATMAP
To demonstrate the effectiveness of detail enhancement in our squeeze-enhanced Axial attention (SEA attention), we ablate our model by removing the detail enhancement.
We visualize the attention heatmaps of the two models in Figure 5.
Without detail enhancement, attention heatmaps from solely SA attention appears to be axial strips while our proposed SEA attention is able to activate the semantic local region accurately, which is particularly significant in the dense prediction task. [Figure]
H. P REDICTION RESULTS
We show the qualitative results and compare with the alternatives on the ADE20K validation set from two different perspectives.
First we compare with a mobile-friendly rival TopFormer Zhang et al.
(2022c) with similar FLOPs and latency in Figure 6.
Besides, we compare with the Transformerbased counterpart SegFormer-B1 Xie et al.
(2021) in Figure 7.
In particular, our SeaFormer-L has lower computation cost than the SegFormer-B1.
As shown in both figures, we demonstrate better segmentation results than both the mobile counterpart and Transformer-based approach.
I L IMITATIONS AND SOCIETAL IMPACT
The mobile-friendly segmentation is deeply related to the industrial application on edge computation platforms, while few academic attempts are made to meet the requirement of the industry.
We test our method on a Qualcomm Snapdragon 865 processor (Fig.1 main paper) and shows superior results to the alternatives.
We believe our work can lead to expected and unexpected innovations in both academia and industry.
However, our system is not perfect yet and hence not fully trustworthy in real-world deployment.
Also, the current system is not exhaustively evaluated and tested due to limited resources.
We focus on mobile semantic segmentation and image classification tasks.
New mobile-friendly method for more downstream tasks and extended to GPU systems will be studied in the future.
[Figure] [Figure]