EDNet: Efficient Disparity Estimation with Cost Volume Combination and Attention-based Spatial Residual

CNN based matching cost computation methods make a great contribution to the stereo matching accuracy. There are two popular approaches for matching cost computation. The first one is using either a layer of 2D [20] or 1D [5] convolutional operations, called correlation layer. Such an inner product between feature vectors is adopted in [21, 22, 6, 7]. Liang $et$ $al.$ [6] build a correlation volume for initial disparity estimation which follows a disparity refinement module by learning through feature consistency. Wang $et$ $al$. [10] make some modification and propose a point-wise correlation volume to preserve fast computation. Another popular method to compute matching cost is to form a 4D volume by concatenating the corresponding features from the opposite stereo images across each disparity level. 3D convolutions are followed to aggregate features and regress disparity. This method can be found first in [1]. Chang $et$ $al$. [2] improve Kendall’s approach [1] by designing a spatial pyramid pooling module [23] so that correspondence estimation can benefit from the image features with rich object context information. Guo $et$ $al$. [3] combine the concatenation volume with the group-wise correlation volume and improve the accuracy with 3D convolutions. The best performance on Scene Flow dataset comes from [24] which introduces the idea of DenseNet [25] to further improve PSMNet [2]. Zhang $et$ $al$. propose GANet [4] with two guided aggregation layers and fifteen 3D convolutions to achieve state-of-the-art performance.

The residual learning concept is proposed by He $et$ $al$. [9] which turns to be an efficient way to train a CNN model and has been adopted by many works. In the stereo matching task, the residual learning strategy is widely used for refining disparity estimation [6, 10, 26, 12]. Pang $et$ $al$. [8] present a cascade residual learning scheme and adopt a two-stage CNN, in which the second stage refines the estimation by producing residual signals. Stucker $et$ $al$. [27] specially build a U-Net [28] based network to enhance the reconstruction by regressing a residual correction. In order to meet the need for real-time inference, [29] takes residual learning strategy to flexibly output disparity estimation according to the requirement of applications. Song $et$ $al$. [11] manage to aggregate edge information for residual learning and thus construct a multi-task network for edge detection and stereo matching.

The architecture of our proposed EDNet is shown in Figure 2. We exploit the structure of DispNetC [5] as the backbone with extensive modifications. For feature extraction, the last left and right feature maps of conv3 from the weights-sharing encoder network are used to form the combined volume which is composed of a squeezed concatenation volume and correlation volume which will be discussed in Section III-B. The detail of our proposed cost volume combination method can be found in the left bottom corner of Figure 2. 2D convolutions are then employed to aggregate the combined volume and regress the disparity. In the decoder part, we follow the coarse-to-fine strategy to refine the disparity progressively. The spatial attention module is applied in order to generate attention-aware residual features, which will be introduced in Section III-C. The stacked hourglass module in PSMNet [2] is used for residual regression but is implemented by 2D convolutions. The attention-based spatial residual module is well illustrated in the right bottom corner of Figure 2. Different from DispNetC [5] which has 6 scales of output, we reduce the disparity prediction to 4 scales, removing the prediction at 1/16 and 1/32 of full resolution.

Previous works simply build a correlation volume [5, 8, 10] or 4D concatenation volume [1, 3, 2] which follows 2D or 3D convolutions for aggregation. However, a single cost volume cannot meet the need of preserving contextual information and feature similarity at the same time. GwcNet [3] improves the performance by combining the group-wise correlation and concatenation volume, but 3D convolutions are required for aggregation which leads to higher memory consumption and more complex computation. To this end, we propose to combine the correlation volume and 4D concatenation volume in a more efficient way aiming at taking advantages of both two cost volumes.

Given a pair of stereo features $\textbf{f}_{L}$ and $\textbf{f}_{R}$, we follow the 1D-correlation in DispNetC [5] to calculate the correspondence at each disparity level $d$. The correlation volume is computed as:

$\displaystyle\textbf{C}_{corr}(d,x,y)=\frac{1}{N}<\textbf{f}_{L}(x-d,y),\textbf{f}_{R}(x,y)>$

(1)

where $<x_{1},x_{2}>$ is the inner product of two feature vectors $x_{1}$ and $x_{2}$, and $N$ is the channel number of input features. The shape of the correlation volume is $N\times D\times H\times W$, where $N$ denotes the batch size, $D$ is the estimated disparity range and the spatial size is $H\times W$. Then we construct the 4D concatenation volume by concatenating the left and right feature maps, i.e.,

$\displaystyle\textbf{C}_{concat}(d,x,y)=Concat\{\textbf{f}_{L}(x-d,y),\textbf{f}_{R}(x,y)\}$

(2)

After obtaining the concatenation volume with the shape $N\times D\times C\times H\times W$, we use three 3D convolutions for aggregation and compress it into 1 channel. The aggregated concatenation volume now has the shape $N\times D\times 1\times H\times W$. It is then squeezed into $N\times D\times H\times W$, the same shape as the correlation volume. The correlation volume and squeezed concatenation volume are finally concatenated to form the combined volume. In this way, both the contextual information and feature similarity measurement are incorporated in the combined volume. Further aggregation can be done by 2D convolutions instead of 3D convolutions, which are more efficient. The final combined volume is formed as:

$\displaystyle\textbf{C}_{comb}(x,y)=Concat\{\textbf{C}_{corr}(x,y),\textbf{C}_{concat}(x,y)\}$

(3)

The normal residual learning method lacks the spatial evidence about where the errors occur. We propose an attention-based spatial residual module to guide the residual learning process to pay more attention to those inaccurate regions across the whole spatial space. According to the estimated disparity $\hat{d}^{s}$ at scale $s$, a synthesized left image $\tilde{I}_{L}^{s}$ can be obtained by warping the right image $I_{R}^{s}$, i.e.,

$\displaystyle\tilde{I}_{L}^{s}(x,y)=I_{R}^{s}(x+\hat{d}^{s}(x,y),y)$

(4)

With the warped left image and target left image, we can get the error $E_{L}^{s}=|\tilde{I}_{L}^{s}-I_{L}^{s}|$. A spatial attention module with 3 layers of 2D convolution which are 1$\times$1, 3$\times$3 and 1$\times$1 respectively is applied. The spatial attention feature map $\textbf{f}_{a}^{s}$ is compressed into one channel followed by the sigmoid function to compute the spatial attention vector whose size is $N\times 1\times H\times W$. The input of error map and color stereo images enables the spatial attention module to learn an attention distribution on blurry object boundaries and mismatched pixels. Akin to FADNet [10], CRL[10], FlowNet2 [15], the input to spatial attention module is the concatenation of the stereo images, error map and estimated disparity map.

We follow DispNetC [5] to preserve both the high-level information and local information by skip-connection. The concatenation of ‘upconvolution’ feature maps from the decoder network and corresponding feature maps from the encoder network is then concatenated with the input of attention module to form the residual feature maps $\textbf{f}_{r}^{s}\in R^{H\times W\times C}$. The final attention-aware residual features $\textbf{f}_{ar}^{s}\in R^{H\times W\times C}$ at scale $s$ are computed by multiplying the attention vector and residual features $\textbf{f}_{r}^{s}$, i.e.,

$\displaystyle\textbf{f}_{ar}^{s}=\textbf{f}_{ar}^{s}\otimes\sigma(\textbf{f}_{a}^{s})$

(5)

where $\sigma(\cdot)$ denotes the sigmoid function. The attention-aware residual features are then input to the stacked hourglass module for residual regression. The stacked hourglass module has the same encoder-decoder structure as PSMNet [2] but is implemented with 2D convolutions.

The attention-based spatial residual module is repeated 3 times as we increase the resolution of the disparity map progressively. Different from the aforementioned works [10, 8], we compute error maps across multiple scales as the error information changes accordingly after the correction. Therefore, we provide the residual learning module with constantly updated error maps across multiple scales instead of a single error map at full resolution.

Instead of building a cascade architecture with residual refinement at the second stage [15, 8, 10], we simply replace the direct disparity estimation with a residual estimation over all scales except the smallest $S$ at which the initial disparity is computed. The multi-scale residual outputs are denoted as $\{r^{s}\}_{s=0}^{S-1}$ where 0 represents the scale of full resolution. For the rest $S$ scales, the estimated disparity at the previous scale is first upsampled to the current scale $\hat{d}_{up}^{s}$ using bilinear interpolation function and then added to the residual for refinement. The final predicted disparity $\hat{d}^{s}$ at scale $s$ is produced as:

$\displaystyle\hat{d}^{s}=\hat{d}_{up}^{s}+r^{s},0\leq s\leq S-1$

(6)

Given the output disparity map at different scales, we adopt the pixel-wise smooth L1 loss to train our EDNet at scale $s$:

$\displaystyle L^{s}(d^{s},\hat{d}^{s})=\frac{1}{N}\sum_{i=1}^{N}{smooth}_{L_{1}}(d_{i}^{s}-\hat{d}_{i}^{s}),$

(7)

where $N$ is the number of pixels of the disparity map, $\hat{d}_{i}^{s}$ is the $i^{th}$ element of the predicted disparity $\hat{d}^{s}$, and $d^{s}$ represents the ground truth disparity. The smooth L1 loss function is:

$\displaystyle{smooth}_{L_{1}}(x)=\begin{cases}0.5x^{2},&\text{if }|x|<1\\ |x|-0.5,&\text{otherwise}.\end{cases}$

(8)

The final loss function is a combination of losses over all scales, i.e.,

$\displaystyle L=\sum_{s=0}^{S}{\lambda}^{s}L^{s}(d^{s},\hat{d}^{s})$

(9)

where ${\lambda}^{s}$ is a scalar for adjusting the loss weight at scale $s$.

Three public datasets are used for training and testing our EDNet. The Scene Flow dataset [5] consists of 39,824 pairs of synthetic stereo RGB images (35,454 pairs for training and 4,370 pairs for testing) with a full resolution of 960$\times$540. Both KITTI 2012 [17] and KITTI 2015 [16] are datasets of real scenes with a full resolution of 1242$\times$375. The ground truth of these two datasets is generated by lidar so that only sparse ground truth is available. We evaluate our model on Scene Flow dataset with the end-point error (EPE), 1-pixel error and 3-pixel error. The end-point error computes the mean disparity error in pixels while the 1-pixel error and 3-pixel error measure the average percentage of the pixel whose EPE is bigger than 1 pixel and 3 pixels respectively. The official metrics (e.g., D1-all) are reported for evaluation on KITTI 2012 and KITTI 2015 datasets.

We implemented our EDNet by PyTorch[30] and trained the model with Adam (momentum=0.9, beta=0.999) as optimizer. For Scene Flow dataset, raw images are randomly cropped to 320$\times$640 as input. Our training is performed on 2 NVIDIA RTX 2080ti GPUs for 70 epochs with a batch size of 8 (4 on each GPU). We follow the training strategy in AANet [12], where the initial learning rate is set to 0.001 and decreased by half every 10 epochs after 20^th epoch. The loss weights are set to ${\lambda}_{0}=1.0,{\lambda}_{1}={\lambda}_{2}=0.8,{\lambda}_{3}=0.6$. The crop size for KITTI is set as 256$\times$512. Due to the insufficient training samples of both KITTI 2012 and KITTI 2015, the pre-trained Scene Flow model is used for fine-tuning on the mixed KITTI 2012 and KITTI 2015 training sets for the first 1000 epochs which follows another 400 epochs of training on two datasets to get the submission result respectively. We use a constant learning rate of 0.0001 for KITTI datasets. Inspired by [31] that searching the correspondence at a coarse scale can be beneficial, especially in low-texture or textureless regions, the loss weights are set as ${\lambda}_{0}=0.6,{\lambda}_{1}={\lambda}_{2}=0.8,{\lambda}_{3}=1.0$. For all datasets, color normalization is taken into use with the mean ([0.485, 0.456, 0.406]) and variation ([0.229, 0.224, 0.225]) of the ImageNet [32] for data pre-processing. The maximum disparity is set as 192 pixels.

To validate the effectiveness of each component in EDNet, we evaluate our model with different configurations on Scene Flow dataset. All the experimental results are obtained after 10 epochs of training. As shown in Table I, removing either of the combined volume or attention-based spatial residual leads to a clear drop of performance.

Cost Volume Combination: As shown in Figure 3, models without the combined volume suffer from inaccurate and discontinuous disparity estimation, especially in low-texture regions. The possible reason might be that a single cost volume is unable to avoid the information loss which leads to less robust feature representations. As shown in Table I, the EPE decreases from EDNet-NS’s 1.07 to EDNet-F’s 1.00. The comparison among EDNet-NRS, EDNet-NRCo and EDNet-NR can validate the effectiveness of our proposed combined volume as well.

Attention-based Spatial Residual: It can be directly analyzed from EDNet-NR and EDNet-NA in Table I that the residual learning module brings significant improvement in terms of accuracy, about 36% decrease of EPE. The comparison between EDNet-NA and EDNet-F in Table I demonstrates that the learning efficiency can be further improved with the help of the attention mechanism. The visualization of EDNet-NA and EDNet-F shown in Figure 3 illustrates that more details like the shaper object boundaries can be recovered with the attention-based spatial residual.

Multi-scale Error Maps: Experiments are conducted to prove that our multi-scale error maps mechanism is of great importance to the residual learning. We remove the error map from the residual features as well as the input of attention module from scale 2 to all scale. As shown in Figure 4, models with error maps at multiple scales can achieve a lower EPE loss with a faster convergence speed. Such performance gain comes at a low cost of extra computation. Multi-scale error maps enable the residual module to learn from the error at the corresponding scale and thus greatly explore the ability of residual learning.

In this subsection, we compare our method with those existing state-of-the-art methods in the aspect of inference speed, memory consumption and accuracy on Scene Flow, KITTI 2015 and KITTI 2012 datasets. We also evaluate the model generalization on Middlebury 2014 dataset [34].

Scene Flow Dataset: As shown in Table II, our proposed EDNet not only outperforms all the competing state-of-the-art methods with the lowest EPE but also runs significantly faster: approximately 7.5$\times$ faster than PSMNet [2], 4$\times$ faster than GwcNet [3], 55$\times$ faster than GANet [4]. Compared with DispNetC [5] and StereoNet [26], EDNet improves the performance by 60% and 40% respectively. Figure 5 gives a visual comparison between our EDNet and other outstanding methods. Sharper object boundaries and more continuous disparity maps can be generated by EDNet, indicating the value of our proposed approach.

KITTI Datasets: We divide the comparison into two groups as shown in Table III and Table IV. First, compared with other top-performing methods, our EDNet still achieves competitive results when evaluating non-occluded pixels while runs considerably faster according to the benchmark. Then we compare our EDNet with previous real-time models. Experimental results from Table III and Table IV show that our work can produce more precise estimation. To stress the efficiency of our proposed method, we compare the computational complexity, memory consumption as well as the inference speed with some popular 3D convolutions based models. Table V comprehensively shows that our model requires less memory consumption: approximately 50% less than PSMNet [2], 40% less than GwcNet [3], 60% less than GANet [4] and 70% less than Bi3D [33]. Moreover, our EDNet has lower computational complexity as our work has the lowest FLOPs (the lower the better) and runs significantly faster: approximately 7$\times$ faster than PSMNet [2], 17$\times$ faster than Bi3D [33] and 45$\times$ faster than GANet [4]. Figure 6 visualizes the disparity and error maps of our method and other state-of-the-art works on KITTI 2015 dataset.

Middlebury 2014 Dataset: We evaluate the model generalization on Middlebury 2014 dataset without additional training. All the models are pretrained on Scene Flow and finetuned on KITTI. As observed from Figure 7, the performance of our EDNet is superior to others, as more smooth and continuous disparity estimation in low-texture regions can be produced. Moreover, our EDNet can preserve better overall object structures and generate sharper edges. The generalization results show that our model has a better adaptation ability to unknown scenes.