Specialize and Fuse: Pyramidal Output Representation for Semantic Segmentation

We adopt the coarse-to-fine pyramid structure to build our pyramidal output format, where a finer level has double resolution than its adjacent coarser level, and all cells (i.e., a patch of pixels) except those in the finest level have exactly four children. Besides, the width and height of an input image should be divisible by those of the coarsest level; otherwise, we resize the input RGB to the nearest divisible.

We denote the index of the pyramid level by $\ell$, where $\ell{=}1$ is the coarsest level and $\ell{=}L$ is the finest level ($L$ levels in total). Let $s_{\ell}$ denote the spatial stride at the pyramid level $\ell$, $D$ the latent dimension of backbone features, and $C$ the number of output classes. As shown in Fig. 2(b) and Fig. 2(c), our model takes a feature tensor $X\in\mathbb{R}^{D\times\frac{H}{s_{L}}\times\frac{W}{s_{L}}}$ from the backbone, and predicts a semantic pyramid $\big{\{}\hat{Y}^{(\ell)}\in\mathbb{R}^{C\times\frac{H}{s_{\ell}}\times\frac{W}{s_{\ell}}}\big{\}}_{\ell=1,\ldots,L}$ and a unity pyramid $\big{\{}\hat{U}^{(\ell)}\in\mathbb{R}^{\frac{H}{s_{\ell}}\times\frac{W}{s_{\ell}}}\big{\}}_{\ell=1,\ldots,L-1}$. The channel dimension of $\hat{U}$ is 1 (binary classification) and is discarded. Note that the output stride of the finest level $s_{L}$ is the same as the output stride of the backbone feature $X$ in this work for simplicity. The final semantic output fused from two pyramids is denoted as $\hat{Y}\in\mathbb{R}^{C\times\frac{H}{s_{L}}\times\frac{W}{s_{L}}}$. The pyramidal ground truths $Y^{(\ell)},U^{(\ell)}$ for training are derived from the ground truth per-pixel semantic labeling $Y$.

At pyramid level $\ell$, each cell is responsible for a patch of $s_{\ell}\times s_{\ell}$ pixels in the original image. A cell can be either an unity cell with the same label for all pixels within the cell, or a mix-cell with more than one labels. In the ground truth unity pyramid $U^{(\ell)}$, positive and negative values indicate unity-cells and mix-cells, respectively. In the ground truth semantic pyramid $Y^{(\ell)}$, for a unity-cell, its ground truth semantic label is defined as the shared label by all covered pixels in the original per-pixel ground truth $Y$, whereas, for a mix-cell, its ground truth semantic label is ill-defined and ignored in computing training loss. Note that a unity-cell at level $\ell$ implies its child cells are also unity-cells at level $\ell+1$. To avoid redundancy, the children unity-cells (referred to as “done by coarser”) are ignored both in computing training loss (Sec. 3.2) and in the fuse-phase (Sec. 3.3).

We first employ a $\operatorname{Conv1x1}$ layer to project the number of channels from backbone’s $D$ to $D_{s}$, producing feature $X^{(L)}$. The simplest way to predict the semantic pyramid $\{\hat{Y}^{(\ell)}\}_{\ell=1,\ldots,L}$ is directly pooling $X^{(L)}$ to the $L$ desired spatial sizes:

where $\operatorname{AvgPool}(\cdot,k)$ is average pooling with kernel size and stride set to $k$. Each $X^{(\ell)}$ is then projected from latent dimension $D_{s}$ to the number of classes $C$ with convolutional layers. We show in our experiment that such a simplest network setting can already achieve promising improvements with the proposed specialize-and-fuse strategy.

Motivated by the recent success of context spreading strategy for semantic segmentation, we further design a coarse-to-fine contextual module for our pyramidal output format. Intuitively, we aggregate the contextual information from coarser pyramid levels to assist the prediction at finer pyramid levels. Fig. 2(c) depicts the proposed coarse-to-fine semantic head, which refines $\{X^{(\ell)}\}_{\ell=1,\ldots,L}$ in Eq. (3) with context feature $\{\theta^{(\ell)}\}_{\ell=0,\ldots,L-1}$ from coarse to fine level:

where the $\operatorname{CA}$ and $\operatorname{CU}$ stand for Context Aggregation module and Context Updating module and will be detailed later. Inspired by ANL’s strategy  [34] to improve efficiency of contextual module, our initial context feature is

where $\operatorname{PyramidPool}$ flatten the 1$\times$1, 3$\times$3, 6$\times$6, and 8$\times$8 features generated by spatial pyramid pooling [15].

Iterating $\ell$ from coarse to fine (from $1$ to $L$), the Context Aggregation module refines $X^{(\ell)}$ using $\theta^{(\ell-1)}$ (while $X^{(1)}$ is refined by the initial context feature $\theta^{(0)}$); the Context Updating module then updates the context feature $\theta^{(\ell-1)}$ with the refined $X^{\prime(\ell)}$, forming the new context feature $\theta^{(\ell)}$ which facilitates finer-level semantic prediction by encapsulating the information from the coarsest to the current level.

Once the coarse-to-fine contextual module generates the refined feature pyramid $\{X^{\prime(\ell)}\}_{\ell=1,\ldots,L}$, the semantic pyramid $\{\hat{Y}^{(\ell)}\}_{\ell=1,\ldots,L}$ are predicted by

where $\operatorname{ConvBlock}^{(\ell)}$ consists of $\operatorname{Conv1x1}$, $\operatorname{BN}$, $\operatorname{ReLU}$, and a final $\operatorname{Conv1x1}$ projecting $D_{s}$ to the number of classes $C$. In below, we detail the Context Aggregation module ($\operatorname{CA}$) and Context Updating module ($\operatorname{CU}$).

The context aggregation module is illustrated in Fig. 2(d). To refine $X^{(\ell)}$, we use attention operation to aggregate coarser-level context encoded in $\theta^{(\ell-1)}$. Specifically, we transform $X^{(\ell)}$ to $X^{(\ell)}_{\text{query}}$, $\theta^{(\ell-1)}$ to $\theta^{(\ell-1)}_{\text{key}},\theta^{(\ell-1)}_{\text{value}}$ by $\operatorname{Conv1x1}$ layers; then we apply

where the $\operatorname{Attention}$ is the attention operation [22], $\operatorname{Conv1x1}_{\text{agg.}}^{(\ell)}$ consisting of $\operatorname{Conv1x1},\operatorname{BN},\operatorname{ReLU}$ projects the concatenated $2D_{s}$ channels back to $D_{s}$.

The context updating module is illustrated in Fig. 2(e). To update the context feature $\theta^{(\ell-1)}$ on the refined feature $X^{\prime(\ell)}$ at level $\ell$, we apply

where the $\operatorname{Conv1x1}_{\text{upd.}}^{(\ell)}$ consisting of $\operatorname{Conv1x1},\operatorname{BN},\operatorname{ReLU}$ projects the concatenated $2D_{s}$ channels back to $D_{s}$. Note that the new context feature $\theta^{(\ell)}$ remains at the same low spatial resolution as $\theta^{(\ell-1)}$, so the overall coarse-to-fine contextual module runs efficiently.

ADE20k is a dataset with diverse scenes containing 35 stuff and 115 thing classes. The training, validation, and test split contains 20K/2K/3K images, respectively. The results on validation set of ADE20K is shown in Table 1. Our method establishes new state-of-the-art results with both ResNet101 and HRNet48 backbone. In addition, we submit our prediction on hold-out test set to ADE20k’s server for official evaluation. The results reported in Table 2 also show our advantage over previous methods.

COCO-Stuff is a challenging dataset with 91 stuff and 80 thing classes. The training and the test sets contain 9K and 1K images, respectively. In Table 3, our approach shows comparable performance to recent state-of-the-art results with ResNet101 backbone and outperforms previous methods with HRNet48 backbone.

Pascal-Context is a widely-used dataset for semantic segmentation. It contains 59 classes and one background class, consisting of 4,998 training and 5,105 test images. Results on Pascal-Context test set is presented in Table 4. With both ResNet101 and HRNet48 backbone, our method achieves state-of-the-art results comparing to previous methods with the same backbone.

As demonstrated in Table 6, our pyramidal “output” representation (the second row) consistently brings improvement over the standard single-level output (the first row) under various model settings. It gains $+1.65$, $+1.05$ and $+2.31$ mIoU improvement without any contextual module, with ANL module [34] and with the proposed coarse-to-fine contextual module, respectively.

Our contextual module is designed to accord with the multi-level essence of the proposed pyramidal output. As shown in Table 6, when predicting the standard output (of a single finest level), our coarse-to-fine contextual module and the ANL [34] gain similar improvement over the baseline: $+1.58$ ($40.42\rightarrow 42.00$) and $+1.60$ ($40.42\rightarrow 42.02$). However, when predicting the proposed pyramidal output, our coarse-to-fine contextual module achieves a more significant $+2.24$ improvement ($42.07\rightarrow 44.31$) than appending a single ANL to the backbone $+1.00$ ($42.07\rightarrow 43.07$). Furthermore, our module also gains more improvement than applying ANL modules at all pyramid level (denoted as ANL-multi) $+1.38$ ($42.07\rightarrow 43.45$), indicating the effectiveness of our design to aggregate contextual information from coarser pyramid levels.

To enforce specialization in training phase, it is important to relabel cells which are “done by coarser” as “don’t care”. In Table 7, we show the results comparing to other relabeling policies. We can see that naively training all semantic pyramid levels to predict full semantic segmentation map do not achieve any improvement. This is reasonable as it ignores the fact that only a small portion of the cells in finer pyramid level would be activated in testing. A simple fix is explicitly relabeling all descendants of a ground-truth unity-cell as “don’t care” during training. By doing so, different semantic pyramid levels now can specialize in the pixels assigned by the oracle. A clear $+1.12$ mIoU improvement is now shown. However, the simple fix still ignores the fact that the unity-cell prediction may produce false negatives where a ground-truth unity-cell is falsely classified as a mix-cell and thus a finer-level semantic prediction is referred in inference phase. As a result, the false negatives might make the train-test distributions inconsistent. Finally, our design of re-labeling the descendants of a true positive unity-cell as “don’t care” gains an extra $+0.53$ mIoU improvement.

One may argue that the improvement from our pyramidal output representation stems from the rich supervision at the multi-level rather than our “specialize and fuse” process. To clarify the contribution, we skip the “fuse” phase at inference and use the semantic output at the finest-level as the final output. In this case, the supervision at multi-level are considered as auxiliary supervision. In Table 8, we can see that the auxiliary supervision indeed improves the performance ($+0.63$). However, the improvement is not comparable to our specialized and fuse process ($+2.31$), which suggests that the main contributor to our superior performance is not multi-level supervision but the specialized and fuse strategy.

This work focuses on the setting of $L{=}4$ pyramid levels of output strides $\{4,8,16,32\}$ for our pyramidal output representation in all experiments. We also experiment with $L{=}2$ levels of output strides $\{4,32\}$ which yields a slightly worse results ($44.31$ vs. $44.20$ mIoU). Therefore, we stick to the $L{=}4$ setting.

In Table 9, we divide pixels into four groups (the four columns, $\ell\in\{1,2,3,4\}$) according to the assignment from our predicted unity pyramid (Sec. 3.3). For each group, we show the performance of the semantic predictions of the four pyramid levels (the four rows, $\ell^{\prime}\in\{1,2,3,4\}$). We can see that the mIoU is degraded if a pixel refers to a pyramid level that does not agree with the assignment by the trained unity pyramid ($\ell^{\prime}\neq\ell$ in each column of Table. 9), indicating that different semantic pyramid levels learn to specialize in predicting the pixels assigned by our predicted unity pyramid.

To demonstrate our intuition that different pyramid levels have their specializations in different classes, in each row of Fig. 3, we show the per-class IoUs predicted by each semantic level. The results suggest that each $\hat{Y}^{(\ell)}$ learns better for different classes in practice. For instance, $\hat{Y}^{(4)}$ performs better at trafficlight, while $\hat{Y}^{(2)}$ is good at mountain. See supplementary material for more visualizations.

We show some qualitative results comparing to the strong baseline, HRNet48-ANL, in Fig. D. See the supplementary material for more examples.