Realtime Global Attention Network for Semantic Segmentation

Zeng et al [1] implemented two FCNs [5] with ResNet-101 [6] backbones to fuse color and depth streams of cluttered scene. This approach yields 83.4% precision at Top-1% percentile affordance proposals. Azpiri et at [7] further refined the Top-1% percentile precision to $\sim$94% by a deep Graph-Convolutional Network backbone, which outperforms FCN (ResNet-101 backbone) by $\sim$2%. Neither of two approaches takes colored images as the only input, consuming much inference time in processing depth information. Shao et at [8] proposed to predict affordance map by sharing a ResNet-50 backbone between colored image stream and depth image stream, and a U-Net [9] to fuse features at different scales. Their approach is annotation-free, but requires the robotic hand attempting to find the most possible points for suction, which is more expensive than deriving direct predictions from monocular images.

Self-Attention mechanism was first proposed in the field of nature language processing [3] for temporal domain. Wang et al [10] adapted the idea to spatial domain, that a feature vector is spatially related to all other features in the feature volume, the features that are highly related will generate strong response, which facilitates network to model the non-local relations across the entire volume.

Li et al [11] designed a fully-convolutional feature pyramid attention module to replace the spatial-pyramid-pooling module [12], and a global-attention-upsample module with which high-level features perform global average pooling [13] as the guidance for low-level features. Woo et al [14] believed that simple channels-wise attention and spatial-wise attention sub-modules can boost representation power of CNNs. Recently, OCNet [15] efficiently aligns a global relation module and a local relation module, dividing and merging feature volumes in different styles. Bello et al [16] proposed to augment standard convolution by attention mechanism. Cao et at [17] proposed a global-context block that benefits from simplified non-local block [10] and squeeze-excitation module [18]. The proposed GAM in this paper is inspired by the criss-cross attention (CCA) mechanism of CCNet [2]. CCNet approximates non-local attention by two cascade CCA modules, each of which correlates all feature vectors aligned horizontally and vertically.

RGANet5 is designed as a single ‘distillation tower’ (see Fig. 2) with 5 temperature levels, each level is composed of one IB. The architecture is expandable such that RGANet$n$ has $n$ cascade IBs. Each IB halves input feature size, then modulates channels to ratio $k=15$ by a standard $3\times 3$ convolution. ESS block (see Fig. 3(a)) is regarded as the backbone for IB, we adopt ESS-3, ESS-6, ESS-12, ESS-24 for IB-1 and IB-2, IB-3, IB-4, IB-5 respectively. Output of GAM (refer to Fig. 3(b)) is directly added to the output feature volume x of ESS block as the residual $\lambda\cdot\textbf{x}$, which formulates the incremental up-sampling layers accordingly. The final output takes the form $(1+\lambda)\cdot\textbf{x}$, where $\lambda$ is the learnable weights volume that has the same size of x, and operator ‘$\cdot$’ signifies element-wise product. Therefore, non-negative activation of GAM artifact is preferred, and Batch-Norm (BN) layers are necessary to restrict its upper bound.

As the synthesis of 5-scales distillation artifacts, highway connections not only populate previous inferences to up-sampling layers accordingly, they also facilitate gradient back-propagation during training phase, especially for a very deep network. The ‘vote & upsample’ (VU) block (see Fig. 2) weights concatenated input feature maps by $1\times 1$ convolution without bias, these weighted features are then $2\times$ upsampled via nearest interpolation. If inferring without last two ESS-3 blocks, RGANet will not yield fine-grained predications. Also, we noticed that IB-1 is directly linked to the last UV4 block, which performs poorly without Decision Unit (DU) shown in Fig. 2. DU consists of feature modulating head to deduct channels from $k+3$ to the number of classes by $1\times 1$ convolution, and one last GAM that yields predictions.

Fig. 3(b) illustrates the pipeline of the GAM. With a batch-size of 1, input feature volume $\bf{x}$ filtered by previous ESS block has a shape of $h\times w\times c$, which is rotated to $c\times h\times w$ query $Q(c,h,w)$ (upper branch, blue color) and $w\times c\times h$ key $K(w,c,h)$ (lower branch, purple color). Query conducts channel-wise attention via depth-wise sliding $w$ kernels $W(c,1)$ shaped $c\times 1$ across $c-h$ view, which results in a $c\times 1\times w$ feature volume $Q^{c}$. Similarly, $h$ horizontal positions are encoded by $1\times h$ depth-wise convolutional kernel $W(1,h)$, mapping into the $c\times 1\times w$ horizontally positional encoding $Q^{h}$. Key is transformed accordingly except the sizes of kernels, we denote its artifacts as channel-wise encoding $K^{c}$ and vertically positional encoding $K^{v}$. The final output $F_{out}$ are then formulated as:

where $\textrm{rot}[\cdot]$ operator signifies rotation; $f\{\cdot\}$ operator performs $1\times 1$ convolution to resize $2c$-channels to $c$-channels with Swish activation function [19], then maps to [0, 1] weights volume via Sigmoid or Softmax functions; ‘$lhs~{}\bigcup~{}rhs$’ operator concatenate $lhs$ and $rhs$ features; Let ‘$*$’ denotes the depth-wise convolution, we know:

Consider Eqn. (1) and Eqn. (2), RGA’s total number of trainable parameters is $2hw+cw+hc$ (exclude BN layer, bias and $1\times 1$ point-wise convolution), while for vanilla convolution with the same kernel height and width, this number becomes $hw^{2}+wh^{2}+cw^{2}+ch^{2}$. In terms of global attention, although matrix multiplication operation only correlates features horizontally and vertically aligned for any feature vector, $Q^{h}$, $Q^{c}$, $K^{v}$ and $K^{c}$ themselves are the artifacts of global depth-wise convolutions, which, as a result, each element in these 4 Queries and Keys bounds other elements with shared weights. We can also treat RGA as the type of ‘learnable global attention’. Furthermore, all affine operations of RGA module are differentiable, we didn’t observe vanishing, or exploding gradients issues during training.

A Dense block [20] has multiple densely-connected Bottlenecks (BNK) modules, each BNK acts in a stack-squeeze manner, and the last one outputs a $k$-channels feature volume.

The ESS-$n$ block, as illustrated in Fig. 3(a), outputs a stacked $(n+1)k$-channels feature volume by $n$ densely-connected BNKs, each contributes $k$-channels of features, and we let $C_{in}=k$. Each channel of stacked feature maps shows one degree of filtered raw input features, such that RGA modules are capable of correlating high-order features to low-level features. RGANet5 adopts ESS-3 for IB-1, ESS-3 for IB-2, ESS-6 for IB-3, ESS-12 for IB-4, and ESS-24 for IB-5. One of our future works is to let RGA module ignore noisy low-level features, and locate more reliable high-order features.

Existing evaluation metrics poorly treat the special case of a prediction map as illustrated in Fig. 4 that, when original predictions cover object-1 (red), but fail to locate object-2 (green) at the left-bottom corner of image, off-the-shelf metrics yields the same results if part of object-1 relocates to the object-2, because this relocation does not affect paranormal statistics of TP, FN, TN and FP. In reality, we want predications to be able to cover more objects, such that a robotic hand would seek-and-pick all objects even though Precision or Recall is yet not favorable enough (e.g., $\sim$15% Recall on object-2 alone). It would be more reasonable to assign higher score for the prediction map that covers 2 objects.

Public-available suction dataset [21] consists of camera intrinsics and pose records, RGB-D images of clutter scenes and their backgrounds, and labels. We only adopted color images and labels w/ a train-split of 1470 images and a test-split of 367 images, each image has a resolution of 480$\times$640. All colored images were normalized to tensors valued between 0 and 1, which were not resized or padded during training and testing.

We implemented AdamW optimizer [22] with AMSGrad [23], compared two weighted loss function during training - focal loss (FLoss) [24] and cross-entropy loss (CELoss). Assume $y$ denotes the prediction, $y^{\prime}$ the one-hot encoded ground-truth, $n$ the total number of classes, $\gamma$ and $\alpha$ are constants, then FLoss and CEloss are denoted as:

where $\sum_{c=1}^{n}\alpha_{c}=1$ and $\alpha_{c}\geq 0$, $\gamma>0$, and all $y\in[0,1]$. Training-set is augmented during training by random hue, flip, rotation, blur, shift etc.

Online testing comes simultaneously during training, which shows calculations of Jaccard Index, precision, and recall of running batches. The Offline testing loads checkpoint, merely evaluates the class corresponds to predicted suction areas.

All experiments were conducted using a GTX1070 laptop, and one Tesla V100 GPU. Network scaler $k$ was set to 15 constantly for a better trade-off between module size and performance. We set constant learning rate to $1.5\times 10^{-4}$, weights decay rate to 0; $\gamma=1.3$, $\alpha_{1}=0.25$ and $\alpha_{2}=0.25$ upon background and negative samples, $\alpha_{3}=0.5$ upon suction areas because of unbalanced proportions; default MGRID parameters $\beta=0.5$, grid intervals $(\delta_{H},\delta_{W})=(12px,12px)$, $C_{m}$ and $F_{m}$ took the same values as illustrated in Fig. 4. The training of RGANet5 lasted for $\sim$2 days.

We implement multiple metrics to evaluate inferential artifacts. Note that for the suction dataset, users only care about feasible regions to adhere. Therefore, only the category that represents predicted suction areas is evaluated, the final scores are computed via averages over entire test-set.

We conducted experiments to compare RGANet5 with several novel semantic segmentation approaches. These approaches can be divided into two groups - one group that has much more parameters/FLOPs that substantially can outperform RGANet, and the other one that has comparable model sizes. As shown in Tab. II, Deeplabv3 [25] and FCN [5] do not rely on attention mechanism, while CCNet [2] has merely two cascade CCA modules that bring in tremendous amount of trainable parameters and operations. Light-weighted RGANet with 6 GAMs, on the other hand, achieves competitive performance (4-7% less) against the best performer on the test-set. Also, RGANet5 outperforms BiSeNetv1 [26], another attention-based realtime approach, by $\sim 6\times$ less parameters and FLOPs.

As illustrated in Tab. III, proposed approach achieves the best Jaccard Index, Dice coefficient and MGRID score when compared with top-tier realtime approaches selected from Cityscape Leader board [33]. Although behaves better in metrics evaluation, RGANet runs relatively slow due to the fact that PyTorch is well-optimized for convolutional neural networks. It is one of our future works to further optimize RGANet for faster and more accurate inference. Readers may refer to Fig. 6 for our qualitative evaluation.