Squeeze-enhanced axial transformer, its layer and methods thereof

A squeeze-enhanced axial transformer for mobile semantic segmentation, including

• a shared stem,
• a context branch,
• a spatial branch,
• a fusion module and
• a light segmentation head, wherein
• the shared stem produces a feature map which is shared by the context branch and the spatial branch;
• the context branch obtains context-rich information from the feature map;
• the spatial branch obtains spatial information in high resolution from the feature map;
• the fusion module incorporates the features in the context branch into the spatial branch, fusing high-level context-rich information into low-level spatial information;
• and
• the light segmentation head receives the feature from the fusion module and output the results;
• wherein the context branch includes three stages, the second and the third one of the three stages each containing a first MobileNet block and a squeeze-enhanced axial transformer (SeaFormer) layers block stacked after the first MobileNet block;
• said SeaFormer layers block comprising one or more SeaFormer layers, wherein said one or more SeaFormer layers each includes a squeeze-enhanced axial attention module and a Feed-Forward Network (FFN);
• wherein the squeeze-enhanced axial attention module comprises a detail enhancement kernel and a squeeze axial attention sub-module;
• the detail enhancement kernel is in multiplication with the squeeze axial attention sub-module to form the squeeze-enhanced axial attention module;
• the detail enhancement kernel comprises a concatenation block, a first convolution and batch normalization block, an activation layer and a second convolution and batch normalization block arranged in series;
• the squeeze axial attention sub-module comprises a horizontal squeeze block and a vertical squeeze block, wherein said horizontal squeeze block and said vertical squeeze block are arranged in parallel, and said horizontal squeeze block is followed by a first multi-head attention block, while said vertical squeeze block is followed by a second multi-head attention block;
• and
• the squeeze axial attention sub-module further comprises a first convolutional layer, wherein the outputs of the two multi-head attention block are broadcasted and then added in element-wise, and then input to the first convolutional layer;
• wherein query key value projections of the input feature are concatenated on the channel dimension in the concatenation block of the detail enhancement kernel, and then passed to the first convolution and batch normalization block;
• and
• wherein the query key value projections of the input feature are taken average on the horizontal direction in the horizontal squeeze block, and taken average on the vertical direction in the vertical squeeze block.