AASEG: ATTENTION AWARE NETWORK FOR REAL TIME SEMANTIC SEGMENTATION

1 INTRODUCTION . Semantic Segmentation is a fundamental problem in computer vision where the goal is to label each and every pixel in the image to its appropriate class . Since it is required to be deployed in real world settings like robots and autonomous vehicles , hence there is a need to balance the speed vs performance tradeoff . Fully Convolutional Network ( FCN ) ( Long et al. , 2015 ) composed of convolutional layers was one of the first works to get strong semantic representation . However , this method was not able to capture boundary information accurately . Atrous convolutions ( Yu and Koltun , 2015 ) at the last several stages of their network was used to give feature maps with strong semantic representation , thus solving the problem with FCN based architectures . However , this comes at at the cost of increased computational complexity . For real world navigation tasks like autonomous driving , there is a need to improve the Frame Per Second ( FPS ) . Chen et al . ( 2017 ) and ( Yu and Koltun , 2015 ) used dilated convolutions to increase the recpetive field while maintaining the number of parameters . SegNet ( Badrinarayanan et al. , 2017 ) utilizes a small network structure and the skip-connected method to achieve improved FPS . ( Mehta et al. , 2018 ) , ( Paszke et al. , 2016 ) and ( Poudel et al. , 2019 ) proposed unique approaches to tackle real time semantic segmentation problem . 2 RELATED WORK . ( Fu et al. , 2019 ) introduces spatial-wise and channelwise attention modules to enhance the recpetive field ( Paszke et al. , 2016 ) trims a a lot of convolution filters to reduce computation . ICNet ( Zhao et al. , 2018a ) proposed an image cascade network using multiresolution branches . ( Iandola et al. , 2016 ) allows the neural network to find the critical channels of the feature map and select the most suitable channels by itself . ESPNet ( Mehta et al. , 2018 ) introduces an efficient spatial pyramid ( ESP ) , which brings great improvement in both speed and performance . Bilateral Segmentation Network ( BiSeNet ) ( Yu et al. , 2018a ) used two parts : Spatial Path ( SP ) is used to get with the loss of spatial information and Context Path ( CP ) for compressing the receptive field . ( Yu et al. , 2020 ) used multi-path framework to combine the low-level details and high-level semantics . ( Li et al. , 2019b ) utilizes a light-weight backbone to speed up its network and a multi scale feature aggregation to improve accuracy . SwiftNet ( Orsic et al. , 2019 ) used lateral connections to restore the prediction resolution while maintaining the speed . ( Lin et al. , 2017 ) uses a multipath refinement network to refine the feature but ignores the global context feature . The speed-Accuracy performance comparison of state of the art methods on the Cityscapes test set is shown in Figure 1 : . We summarize our main contributions as follows : • We propose a novel Attention Aware Network ( AASeg ) for real time semantic segmentation . • Our network is comprised of three parts : Spatial Attention ( SA ) module to capture the spatial dependencies of the feature maps , Channel Attention ( CA ) module to extract high level semantic information and Multi Scale Context ( MSC ) module to learn the information flow between feature maps of consecutive levels . • Detailed experiments and analysis indicate the efficacy of our proposed network in not only improving the performance but also FPS . We achieve results on par with previous state of the art using Cityscapes , Camvid and on ADE20K datasets . 3 BACKGROUND . 3.1 SPATIAL INFORMATION . The spatial information of the image is important to predict the detailed output for semantic segmentation . Modern existing approaches uses encoding of spatial information . DUC ( Wang et al. , 2018 ) , PSPNet ( Zhao et al. , 2017 ) , DeepLab v2 ( Chen et al. , 2017 ) use the dilated convolution to preserve the spatial size of the feature map . 3.2 CONTEXT INFORMATION . Semantic segmentation requires context information to generate a high-quality result . Most of the used methods enlarge the receptive field or fuse different contextual information . ( Chen et al. , 2017 ) , ( Wang et al. , 2018 ) and ( Yu and Koltun , 2015 ) use different dilation rates in convolution layers to capture different scale contextual information . In ( Chen et al. , 2017 ) , an “ ASPP ” module is used to capture context information of different receptive field . PSPNet ( Zhao et al. , 2017 ) applies a “ PSP ” module which contains several different scales of average pooling layers . 3.3 ATTENTION MECHANISM . ( Hu et al. , 2018 ) applied channel attention for image recognition and achieve the state-of-the-art . ( Yu et al. , 2018b ) proposed a network that learns the global context as attention and revise the features . Multi scale network along with a custom attention module was proposed by ( Sagar and Soundrapandiyan , 2020 ) for semantic segmentation . 3.4 FEATURE FUSION . The spatial information captured by the Spatial Path consists of rich detailed information . The output feature of the Context Path is made up of contextual information . Feature fusion was used to achieve state of the art results on image classification , object detection and instance segmentation using DMSANet ( Sagar , 2021b ) . 4 PROPOSED METHOD . 4.1 DATASET . The following datasets have been used to benchmark our results : 1 . Cityscapes It is used for urban street segmentation . The 5000 annotated images are used in our experiments which are divided into 2975 , 500 and 1525 images for training , validation , and testing respectively . 2 . ADE20K This dataset contains labels of 150 object categories . The dataset includes 20k,2k and 3k images for training , validation and testing respectively . 3 . CamVid This dataset is used for semantic segmentation for autonomous driving scenarios . It is composed of 701 densely annotated images . 4.2 NETWORK ARCHITECTURE . The fundamental goal in semantic segmentation is to map an RGB image X ∈ RH×W×3 to a semantic map Y ∈ RH×W×C with the same spatial resolution H ×W , where C is the number of classes of objects present in image . The input image X is converted to a set of feature maps Fl where l=1 , ... ,3 from each network stage , where Fl ∈ RHlWlCl is a Cl-dimensional feature map . Our network dosen ’ t use any backbone to extract features from the input image unlike many previous architectures . The input image is first passed through a block comprising of convolutional , batch normalization and ReLU activation function . We express a convolution layer Wn ( x ) as follows : Wn ( x ) = Wn×n x+ b ( 1 ) where represents the convolution operator , Wn×n represents the n × n convolutional kernel , x represents the input data and b represents the bias vector . 4.3 MULTI SCALE CONTEXT MODULE . The feature map after passing through convolutional block is split to three different parts with 1×1 convolution , 3×3 convolution and 5×5 convolution respectively . The individual feature maps are then fused together . The output fused feature map is first convolved with a 1 × 1 convolution to reduce the number of channels from 2048 to 256 . The feature map produced is of size H ×W × nc , where H , W are the height and width of the feature map , and nc denotes the number of channels . The input feature map is convolved with dilated convolution layers with increasing dilation rates of 3 , 6 and 12 . The dilated convolution layer input at every stage is formed by concatenating the input feature map with the outputs from previous convolutions . At the final step , the outputs from the three dilated convolutions are concatenated with the input feature map . 4.4 SPATIAL ATTENTION MODULE . The spatial attention module is used for capturing the spatial dependencies of the feature maps . The spatial attention ( SA ) module in our network is defined below : fSA ( x ) = fsigmoid ( W2 ( fReLU ( W1 ( x ) ) ) ) ( 2 ) where W1 and W2 denotes the first and second 1× 1 convolution layer respectively , x denotes the input data , fSigmoid denotes the sigmoid function , fReLU denotes the ReLu activation function . The spatial attention module used in this work is shown in Figure 2 : 4.5 CHANNEL ATTENTION MODULE . The channel attention module is used for extracting high level multi-scale semantic information . The channel attention ( CA ) module in our network is defined below : fCA ( x ) = fsigmoid ( W2 ( fReLU ( W1f 1 AvgPool ( x ) ) ) ) ( 3 ) where W1 and W2 denotes the first and second 1 × 1 convolution layer , x denotes the input data . f1AvgPool denotes the global average pooling function , fSigmoid denotes the Sigmoid function , fReLU denotes ReLU activation function . The channel attention module used in this work is shown in Figure 3 : 4.6 AGGREGATION . We denote the concatenation operation as follows : xconcat = x1 ⊕ x2 ⊕ x3 ( 4 ) where ⊕ represents the concatenation operator and x1 , x2 and x3 represents the features of the two branches . The AASeg module can be denoted as follow : xAASeg = ( ( fSA ( xconcat ) ⊗ xconcat ) ⊕ ( fCA ( xconcat ) ⊗ xconcat ) ⊕ ( fMSC ( xconcat ) ⊗ xconcat ) ) ( 5 ) where ⊕ represents the concatenation operator , fCA represents the channel attention module mentioned in Equation 2 , fSA represents the spatial attention module , fMSC represents the multi scale attention and xconcat represents the combined feature . we use ConvXi to denote the operations of ith block . Therefore , the output of ith block is calculated as follows : xi = ConvXi ( xi−1 , ki ) ( 6 ) where ConvX includes one of each convolutional layer , batch normalization layer and ReLU activation layer , ki is the kernel size of convolutional layer , xi−1 and xi are the input and output of ith block . Fusion operation is used to combine high-level features with low-level features using Equation 7 : xoutput = F ( x1 , x2 , . . . , xn ) ( 7 ) The overall structure of our proposed AASeg network is shown in Figure 4 . 4.7 LOSS FUNCTIONS . We use cross-entropy loss function to weigh the difference between the forward propagation result of the network and the ground truth of the samples . The cross-entropy loss is calculated as defined in Equation 8 : Lce = 1 N N∑ n=1 [ yn log ŷn + ( 1− yn ) log ( 1− ŷn ) ] ( 8 ) where N denotes the total number of samples , yn denotes the probability that the forward propagation result is true , and 1-yn denotes the probability that the forward propagation result is false . We also use the auxiliary supervision loss Laux to improve the model performance thus making it easier to optimize . The auxiliary loss can be defined as : Laux = − 1 BN B∑ i=1 N∑ j=1 K∑ k=1 I ( gij = k ) log exp ( pij , k ) ∑K m=1 exp ( pij , m ) ( 9 ) I ( gij = k ) = { 1 , gij = k 0 , otherwise ( 10 ) where B is the mini batch size , N is the number of pixels in every batch ; K is the number of categories ; pijk is the prediction of the jth pixel in the ith sample for the kth class , I ( gij = k ) is a function which is defined in Equation 10 . The class attention loss Lcls from channel attention module is also used . The class attention loss is defined as follows : Lcls = − 1 BN B∑ i=1 N∑ j=1 K∑ k=1 I ( gij = k ) log exp ( aij , k ) ∑K m=1 exp ( aij , m ) ( 11 ) where aijk is the value generated of the class attention map of the jth pixel in the ith sample for the kth class . We combine the three terms to balance final loss term as follows : L = λ1Lce + λ2Lcls + λ3Laux ( 12 ) where λ1 , λ2 and λ3 are set as 1 , 0.5 and 0.5 to balance the loss .

Summary: This paper proposes a real-time semantic segmentation. To save computation, there is no heavy backbone to extract features but just simple block containing conv/relu/bn. Then there three modules applied in parallel on top of the image features: multi-scale context module (conv with different dilation rate), spatial attention module and channel attention module. Experiments were done on three datasets and the proposed method can achieve a reasonable performance with higher FPS than previous methods.