Semi-Global Shape-aware Network

Self-attention [33] was initially applied in machine translation area. Non-local Networks [35] bridged self-attention modules in machine translation to non-local filtering operations in computer vision. Many methods improved weight functions in self-attention module to learn discriminative feature representation. Non-local Networks [35] discussed four choices of weight functions, i.e. Gaussian, Embedded Gaussian, Dot product and Concatenation. Considering the computational complexity $\mathcal{O}(N^{2})$ of Non-local Networks, CCNet [13] adopted recurrent sparse criss-cross attention modules to substitute the dense attention in Non-local Networks. By two consecutive criss-cross attention modules, every node could collect contextual information from all nodes in feature maps. This way reduced the computational complexity $\mathcal{O}(N^{2})$ for Non-local Networks to $\mathcal{O}(N\sqrt{N})$, therefore it was more efficient, while object shapes could be further considered. Song et al. [29] first built minimum spanning trees (MST) in low-level feature maps and then computed feature similarity in high-level semantics to preserve object structures. However extra time and computations to build the minimum spanning trees were needed. In this work, we consider both spatial distances and feature similarity when capturing long-range dependencies for higher accuracies meanwhile maintaining a low computational complexity.

Semantic segmentation is an essential and challenging task in computer vision community. The methods based on Convolution Neural Networks (CNNs) have made significant achievements in the past few years. Long et al. [17] applied fully convolutional networks (FCNs) in semantic segmentation. Later, researchers found that FCNs were limited by receptive fields due to fixed geometric structures. The works of U-net [27] and DeepLabv3+ [7] used encoder-decoder structures to combine high-level and low-level features for semantic segmentation tasks. Chen et al. [5] adopted atrous convolutions which could effectively enlarge receptive fields when aggregating contextual information. Furthermore, they also proposed atrous spatial pyramid pooling (ASPP) to complete the segmentation task at multiple scales. Zhao et al. [42] exploited global context information by pyramid pooling modules to achieve better performance. PSANet [43] relaxed the local neighborhood constraint through a self-adaptively learned attention mask. GCN [22] found that large convolutional kernels were also important when performing a dense per-pixel prediction task and proposed Global Convolutional Network (GCN). Unlike the previous methods, we aggregate global context information in a self-attention manner.

Traditionally, bag-of-visual-words [24], VLAD [2] and Fisher vector [23] are usually used to aggregate a set of handcrafted local features [18] into a single global vector to represent an image. Recently, many methods attempt to replace handcrafted ones by learned counterparts and then aggregate learned features with these similar techniques as the traditional ones [1, 16, 10]. Some studies show that directly using a pooling operation [26] to substitute the aggregation process can get comparable performances. We follow [20] to name these methods with a pooling operation as global single-pass methods for they do not separate extraction and aggregation steps explicitly. Radenović et al. [26] proposed GeM pooling which generalized average and max pooling and got excellent results. Based on GeM, SOLAR [20] employed second-order similarity and attention for image retrieval and obtained significant performance improvements. In this paper, we also explore the proposed SGSNet components for image retrieval.

First, a given feature map $I$ in CNNs can be seen as a connected, undirected graph $G=(N,E)$. The nodes $N$ are all spatial positions in $I$ and the edges $E$ with weights $\omega$ are all connections between two neighbor nodes. Next, we define a weight function between different nodes. Let us consider a simple case: given a pair of neighbor nodes $u$ and $v$, the weight between $u$ and $v$ is defined as

$$\omega(u,v)=\omega(v,u)=exp(-d(u,v)),$$

(1)

where $d(u,v)$ is Euclidean distance between feature vectors of $u$ and $v$. When nodes $u$ and $v$ are not neighbors, the weight function is given by

$$\Omega(u,v)=\prod_{\mathclap{(v_{i},v_{j})\in P_{u,v}}}\omega(v_{i},v_{j})=exp(-\sum_{(v_{i},v_{j})\in P_{u,v}}d(v_{i},v_{j})),$$

(2)

where $v_{i}$ and $v_{j}$ are neighbor nodes in the shortest path $P$ between $u$ and $v$. We denote the feature map after aggregating contextual information as $I^{\prime}$. The aggregation function which takes feature similarity and geometric proximity into account simultaneously is:

$$I^{\prime}_{u}=\dfrac{1}{S_{u}}\sum_{\forall v\in\mathcal{R}}\Omega(u,v)f(I_{v})+I_{u},$$

(3)

where $I^{\prime}_{u}$ and $I_{u}$ denotes the feature vector at $u$ in $I^{\prime}$ and $I$ respectively, $\mathcal{R}$ is all nodes in the graph $G$, the function $f(\cdot)$ means feature transformation and the weight $\Omega(u,v)$ is normalized by $S_{u}=\sum_{\forall v\in\mathcal{R}}\Omega(u,v)$.

In the above aggregation process, a problem is that there may be a lot of paths between arbitrary nodes $u$ and $v$ and it is time-consuming to find the shortest path for all pairs of nodes. In [37, 29], the authors adopt a spanning tree to remove ”unwanted” edges in the four-connected planner graph $G$, thus all the nodes are connected by a minimum spanning tree. The minimum spanning tree can enlarge the geometric proximity between two neighbor nodes if these two nodes are dissimilar in appearance. As a result, low support weights will be assigned between these two nodes which causes less robustness to textures [37]. Besides, establishing and traversing a minimum spanning tree needs to pay extra time and computations. Inspired by [12, 13], we adopt a semi-global manner to overcome this difficulty by considering the supports from nodes in horizontal and vertical directions in a single block, and then add multiple blocks hierarchically for capturing full-image dependencies. The details are described in Sec. 3.2.

Due to the difficulty of minimizing matching costs in a 2D image, SGM [12] utilizes a semi-global approach to aggregate the matching costs along multiple 1D directions. CCNet [13] just collects contextual information in a criss-cross path. In this paper, we establish a semi-global block, as shown in Fig. 2 (a). Given an input feature map $I\in\mathcal{R}^{C\times W\times H}$ with channels $C$, width $W$ and height $H$, the block first utilizes a 1 $\times$ 1 convolution $\lambda$ on $I$ to reduce $C$. The resulting feature map is denoted as $\Lambda\in\mathcal{R}^{C^{\prime}\times W\times H}$ with $C^{\prime}<C$. We compute an attention map $A\in\mathcal{R}^{(H+W-1)\times W\times H}$ by a weight function on $\Lambda$. Channels at a position of $A$ correspond to the weights between this position and other positions in the same column or row respectively. The block also applies another 1 $\times$ 1 convolution $\psi$ on $I$ and the resulting feature map $\Psi$ has the same size as $I$. For each position $u$ in the spatial dimension of $\Psi$, we can obtain a feature vector $\Psi_{u}\in\mathcal{R}^{C}$. We collect the feature vectors of all positions in horizontal and vertical directions of $u$ on $\Psi$, resulting in a matrix $\Theta_{u}\in\mathcal{R}^{(H+W-1)\times C}$. We denote the output feature map of the block as $I^{\prime}\in\mathcal{R}^{C\times W\times H}$, and the feature vector at position $u$ on $I^{\prime}$ by (3) is:

$$I^{\prime}_{u}=\sum_{d\in(0,H+W-1)}A_{u_{d}}\Theta_{u_{d}}+I_{u},$$

(4)

where $d$ is the index of channels of $A$, $A_{u_{d}}$ denotes the value at the $d^{th}$ channel in position $u$ on $A$. $\Theta_{u_{d}}$ denotes the $d^{th}$ feature vector in $\Theta_{u}$, and $I_{u}$ denotes the feature vector at position $u$ on $I$.

It is worth noting that the weight function is quite different from [13]. The aggregation process is shown in Fig. 2 (b), where the weight between nodes $u$ and $v$ depends on all nodes in the path connecting $u$ and $v$. The path is limited in horizontal and vertical directions. Actually, our weight function (2) in the aggregation process considers feature similarity and geometric proximity simultaneously, as described in Sec 3.1. Therefore, our method can preserve object shapes more clearly compared with [13], which will be shown in Sec. 4.1.4. In order to adjust feature similarity between nodes $u$ and $v$, we further add learnable parameters $\alpha$ and $\beta$ in (2):

$$\left\{\begin{array}[]{lr}\Omega(u,v)=exp(-\dfrac{D(u,v)}{\alpha})\ \ \ for\ horizontal&\\ \Omega(u,v)=exp(-\dfrac{D(u,v)}{\beta})\ \ \ for\ vertical,\end{array}\right.$$

(5)

where $D(u,v)=\sum_{(v_{i},v_{j})\in P_{u,v}}d(v_{i},v_{j})$, $v_{i}$ and $v_{j}$ are neighbor nodes in horizontal or vertical path $P$ connecting $u$ and $v$. Considering that width and height of feature maps may be different, we use $\alpha$ and $\beta$ for horizontal and vertical directions respectively. If $\alpha$ ($\beta$) is smaller, feature similarity plays a more important role in (5), and vice versa. Therefore, the learnable parameter $\alpha$ ($\beta$) can adjust the relation of feature similarity and geometric proximity adaptively.

A single semi-global block can only aggregate contextual information in the same row and column of a central node but ignores other directions. We address this problem by multiple hierarchical semi-global blocks. In the first hierarchical level, a feature map $I$ is input into a semi-global block, producing a feature map $I^{\prime}$. Then in the second hierarchical level, another semi-global block takes $I^{\prime}$ as the input and outputs a feature map $I^{\prime\prime}$. Thus, every node in $I^{\prime\prime}$ can aggregate contextual information from all nodes in $I$. Specifically, we denote the attention maps in the first and second hierarchical levels as $A$ and $A^{\prime}$. The process of spreading contextual information in node $v$ (in purple) to node $u$ which is not in the same row and column as $v$ is illustrated in Fig. 3. In the first semi-global block, only nodes in the same row and column can receive the contextual information from $v$, which can be seen at $A$ in Fig. 3. In the second semi-global block, $u$ receives contextual information from $v_{1}$ or $v_{4}$ which have aggregated contextual information from $v$ in the first semi-global block, as shown at $A^{\prime}$ in Fig. 3. Thus, $u$ obtains contextual information of $v$ through the above process.

More generally, every node can capture full-image contextual information by the hierarchical semi-global blocks, which can enhance the capability of a single semi-global block for modeling long-range dependencies [13].

For a clearer explanation, we ignore 1 $\times$ 1 convolution $\lambda$ and $\psi$ in Fig. 2, which does not influence our conclusion. In a semi-global block, one node is supposed to aggregate contextual information of those in vertical and horizontal directions. According to (4), the total computational complexity of a single semi-global block is $\mathcal{O}(N\sqrt{N})$ by the brute force implementation. Specifically, given an input feature map $I$ with $N=W\times H$ nodes, every node needs to aggregate contextual information from other nodes in the same row and column. Thus, we need $(H+W-1)$ computation times of similarity for every node. The total computation is $\mathcal{O}((H\times W)\times(H+W-1))$, which is unfavorable for practical applications.

Noticing that there are extensive repeated computations in the brute force implementation, we propose an optimized alternative which reduces the above computational complexity of a single semi-global block to $\mathcal{O}(N)$. We innovatively regard each of rows and columns on the given feature map as a binary tree and complete interactions between any two nodes on the binary tree in linear time $\mathcal{O}(2\times(H-1))$ or $\mathcal{O}(2\times(W-1))$. Thus total computational complexity is $\mathcal{O}(H\times 2\times(W-1)+W\times 2\times(H-1))=\mathcal{O}(4\times W\times H-2\times(W+H))$ for $W$ columns and $H$ rows. Specifically, let us consider operations on one of the rows or columns, as shown in Fig. 4. We select one node arbitrarily in the given row or column as a root node to establish a binary tree. Except the root node and leaf nodes, every node in the binary tree has a parent node and a child node. Every node is supposed to capture contextual information of others in this binary tree. The process of capturing contextual information is split into aggregation and updating steps:

• •

  We aggregate contextual information from leaf nodes towards the root node recursively according to

  $$A(I_{u})=\left\{\begin{array}[]{lr}I_{u}\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ u\ is\ a\ leaf\ node&\\ I_{u}+\sum\limits_{P(v)=u}\Omega(u,v)A(I_{v})\ \ \ \ \ \ \ others,\end{array}\right.$$

  (6)

  where $I_{u}$ is the input feature vector at node $u$, $P(v)=u$ means $u$ is the parent node of $v$. $\Omega(u,v)$ is the defined weight function between $u$ and $v$ in (5).

• •

  We update features from the root node towards leaf nodes recursively by

  $$U(I_{u})=\left\{\begin{array}[]{lr}A(I_{u})\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ u\ is\ a\ root\ node&\\ \Omega(P(u),u)U(I_{P(u)})&\\ +(1-\Omega^{2}(u,P(u)))A(I_{u})\ \ \ \ \ \ \ \ others,\end{array}\right.$$

  (7)

  where $A(I_{u})$ is computed by (6).

The computational complexity of the aggregation and that of updating steps are both $\mathcal{O}(H-1)$ for columns or $\mathcal{O}(W-1)$ for rows. Thus the total computational complexity is $\mathcal{O}(2\times(H-1))$ or $\mathcal{O}(2\times(W-1))$ for a single binary tree.

Operations on other columns and rows are similar. We do not perform updating steps until the aggregation steps are completed for all binary trees. Thus, the aggregation and updating steps of each binary tree are independent of others, which can be parallel accelerated on GPUs. The output feature map is denoted as $I^{\prime}$. A node $u$ is included in two binary trees ($T_{c},T_{r}$), as shown in Fig. 4. We add the updating results of $u$ in $T_{c}$ and $T_{r}$ to $I_{u}$:

$$I^{\prime}_{u}=T_{c}(U(I_{u}))+T_{r}(U(I_{u}))+I_{u},$$

(8)

where $I^{\prime}_{u}$ is the feature vector at $u$ in $I^{\prime}$, and $I^{\prime}_{u}$ has contained contextual information in horizontal and vertical directions of $u$ in $I$.

We adopt Cityscapes [8] benchmark for semantic segmentation and report the metrics of Mean IoU (mean of class-wise intersection over union). Cityscapes is a semantic segmentation benchmark focusing on urban street scenes. This benchmark contains 5,000 finely annotated images which are divided into 2975, 500, 1525 images as training, validation and testing set. A larger set of 20,000 coarsely annotated images are also provided for supporting methods that exploit large volumes of weakly-labeled data. The dataset also defines 30 visual classes, of which 19 classes are used in our experiments.

Based on the proposed linear time algorithm, SGSNet can be plugged into any deep neural networks conveniently for capturing long-range dependencies. We adopt ResNet-101 [11] (pre-trained on ImageNet) as our backbone with minor changes. The last two down-sampling layers are dropped and dilated convolutions are embedded into subsequent convolutional layers similar to [13].

We use the mini-batch stochastic gradient descent (SGD) with a momentum of 0.9 and a weight decay of 0.0001. A poly learning rate schedule with an initial learning rate of 0.01 and power of 0.9 is employed. We also augment the training set by randomly scaling (0.75 to 2.0 $\times$) and then crop out patches with 769 $\times$ 769 pixels as the network input resolution. The learnable parameter $\alpha$ and $\beta$ are both initialized with 1. The number of channels in $\Lambda$ is one eighth of the number of channels in $I$ and we share the parameters of the hierarchical semi-global blocks.