Group-wise Correlation Stereo Network

In plain words

Two cameras side by side see the same scene from slightly different places. How far away a point is can be read from how much it shifts between the two images. Finding those shifts is the hard part.

Many methods compare each point in the left image with candidate points in the right image, shifted by every possible amount, and store the scores in a large three-dimensional table called a cost volume. A neural network then cleans up that table and reads a depth value from it.

This paper changes how the table is filled. Instead of comparing all the features at once, it cuts them into groups of channels, scores each group separately, and stacks the results. Each group gives a different kind of matching evidence, so the network has more to work with and does not have to learn the matching rule from scratch. The authors also modify the network that refines the table, which makes it cheaper to run at inference. They report better accuracy than earlier stereo networks on the Scene Flow and KITTI benchmarks.

TL;DR

• Input: a rectified left and right image pair; output: a disparity map, the horizontal pixel shift per pixel §1.
• Core change: group-wise correlation. The feature channels are split into 40 groups, and each group is correlated at every disparity, instead of one correlation or one concatenation §3.2, eq. 3, §4.3.
• Refinement: a 3D stacked hourglass network with 1×1×1 shortcut convolutions; auxiliary outputs are dropped at inference §3.3, fig. 2.
• Author-reported: better than the compared methods on Scene Flow, KITTI 2012 and KITTI 2015. KITTI 2015 test D1-all is 2.11% against 2.32% for PSMNet §4.6, tab. 4.

Problem

Stereo depth needs a matching score for every left pixel at every candidate disparity. A full correlation gives only one channel per disparity and throws away much of the feature information §2.2, §3.2. Concatenating the two feature maps keeps the information, but the aggregation network must then learn the similarity function from scratch, which costs parameters and computation §2.2, §3.2. Methods that aggregate hand-crafted matching costs cannot be trained end to end §2.2. The paper aims for matching features that keep accuracy when the 3D aggregation network is made smaller §3.2, §4.5.

Method

The pipeline runs in four stages: extract per-pixel features, build the matching volume, refine it in 3D, and read out disparities §3.1, fig. 1.

1. Unary features. A ResNet-like extractor from PSMNet, with half dilation and without its spatial pyramid pooling §3.1. The outputs of conv2, conv3 and conv4 are concatenated into 320 channels at 1/4 of the input resolution §3.1, tab. 1.
2. Cost volume. Two parts are built and, in the full model, concatenated. The group-wise volume has N_g = 40 channels, one per group [§3.2, tab. 1, §4.3]. The concatenation volume uses features compressed to 12 channels per side by two convolutions, giving 24 channels §3.1, tab. 1. Together they give 64 channels, the GwcNet-gc setting tab. 1.
3. 3D aggregation. A pre-hourglass module of four 3D convolutions with batch normalisation and ReLU and a residual add, followed by three stacked 3D hourglass networks §3.1, §3.3, tab. 1, fig. 2.
4. Output modules. After the pre-hourglass module and after each hourglass, a 3D convolution produces a one-channel volume. It is upsampled to the full disparity range and turned into probabilities with softmax over disparity. The disparity is the expectation over disparity levels (soft argmin) [§3.4, eq. 4, tab. 1].
5. Loss. A weighted sum of smooth L1 losses over the four outputs [§3.4, eq. 5, eq. 6].

Group-wise correlation §3.2, eq. 3. The channels are split evenly into N_g groups. The g-th group of the left features is correlated with the g-th group of the right features shifted by d, and the inner product is divided by the number of channels per group:

C_{gwc}(d,x,y,g)=\frac{1}{N_c/N_g}\left\langle \mathbf{f}_l^g(x,y),\, \mathbf{f}_r^g(x-d,y) \right\rangle eq. 3

Here N_c is the channel count and f^g is group g. The volume has shape [D_max/4, H/4, W/4, N_g] §3.2. With N_g = 1 it reduces to full correlation §3.2.

Soft argmin §3.4, eq. 4, where p_k is the probability of disparity level k:

\widetilde{d}=\sum_{k=0}^{D_{max}-1} k \cdot p_k eq. 4

Loss §3.4, eq. 5, eq. 6. Each output i is weighted by λ_i and compared with the ground truth d* using smooth L1:

L=\sum_{i=0}^{3}\lambda_i \, \text{Smooth}_{L_1}(\widetilde{d}_i-d^*) eq. 5

\text{Smooth}_{L_1}(x)=\begin{cases} 0.5x^2, & |x|<1 \\ |x|-0.5, & \text{otherwise} \end{cases} eq. 6

Changes to the hourglass §3.3, fig. 2, tab. 1:

• An extra auxiliary output module on the pre-hourglass features §3.3.
• The residual connections between output modules are removed, so output modules 0 to 2 can be skipped at inference §3.3, fig. 2.
• 1×1×1 3D convolutions on the shortcut connections inside each hourglass. The paper says these add little computation §3.3, fig. 2, tab. 1.

Group-wise volume, as the code builds it:

# Group-wise correlation volume  [§3.2, eq. 3]
# Input:  left, right features (B, C, H, W), C = 320, at 1/4 resolution
# Output: volume (B, N_g, D, H, W), with D = D_max / 4

def group_wise_volume(left, right, num_groups, max_disp):
    # Channels per group, e.g. 320 / 40 = 8  [§4.3]
    group_size = C // num_groups
    volume = zeros(B, num_groups, max_disp, H, W)
    for d in range(max_disp):
        # Pixel x in the left image pairs with x - d in the right  [eq. 3]
        # Positions with x < d stay zero  [code: models/submodule.py#L53-L63]
        shifted_right = shift_right(right, d)
        # Element-wise product, then mean over each group  [eq. 3]
        product = left * shifted_right
        volume[:, :, d] = mean_per_group(product, num_groups)
    return volume

Network forward pass, with the concatenation volume enabled §3, fig. 1, tab. 1:

# GwcNet forward pass  [§3, fig. 1, tab. 1]
# Input:  left, right images (B, 3, H, W)
# Output: one disparity map per output module, each (B, H, W)

def forward(left, right):
    # 1. Unary features: conv2, conv3 and conv4 maps concatenated  [§3.1]
    left_feats = unary_features(left)      # (B, 320, H/4, W/4)
    right_feats = unary_features(right)    # (B, 320, H/4, W/4)

    # 2. Group-wise correlation volume with 40 groups  [§3.2, eq. 3]
    volume = group_wise_volume(left_feats, right_feats, 40, D_max // 4)

    # 3. Concatenation volume from 12-channel features  [§3.1, tab. 1]
    small_left = compress_to_12(left_feats)
    small_right = compress_to_12(right_feats)
    volume = concat(volume, concat_volume(small_left, small_right))
    # (B, 64, D/4, H/4, W/4)

    # 4. Pre-hourglass module with residual add  [§3.1, tab. 1]
    cost0 = pre_hourglass(volume)

    # 5. Three stacked hourglasses, each fed by the previous one  [§3.3]
    cost1 = hourglass(cost0)
    cost2 = hourglass(cost1)
    cost3 = hourglass(cost2)

    # 6. Output module per stage: score, upsample, softmax, soft argmin
    #    [§3.4, eq. 4]
    disparities = []
    for cost in [cost0, cost1, cost2, cost3]:
        scores = output_module(cost)              # (B, 1, D/4, H/4, W/4)
        scores = upsample(scores, D_max, H, W)    # (B, 1, D, H, W)
        probs = softmax(scores, over_disparity)   # (B, D, H, W)
        disparities.append(soft_argmin(probs))    # [eq. 4]

    # At inference only the last output is returned  [§3.3]
    return disparities

One 3D hourglass §3.3, tab. 1, fig. 2:

# One hourglass with 1x1x1 shortcuts  [§3.3, tab. 1, fig. 2]
# Input: x (B, C, D, H, W), with C = 32  [tab. 1]
# Each 3D conv is followed by batch norm; ReLU where tab. 1 says so

def hourglass(x):
    # Encoder: stride-2 convolutions to 2C, then 4C channels  [tab. 1]
    enc1 = relu(conv3d(x, 2 * C, stride=2))
    enc1_b = relu(conv3d(enc1, 2 * C))
    enc2 = relu(conv3d(enc1_b, 4 * C, stride=2))
    enc2_b = relu(conv3d(enc2, 4 * C))
    # Decoder: transposed convs plus 1x1x1 shortcuts  [tab. 1]
    dec1 = relu(deconv3d(enc2_b, 2 * C) + shortcut(enc1_b))
    return relu(deconv3d(dec1, C) + shortcut(x))

Training loss:

# Training loss  [§3.4, eq. 5, eq. 6]
# Input:  four predicted disparity maps, ground truth d_star (B, H, W)
# Output: scalar loss

def loss(disparities, d_star):
    # Weights lambda_0 to lambda_3  [§4.2]
    weights = [0.5, 0.5, 0.7, 1.0]
    # Not in the paper: only pixels with 0 < d_star < D_max count
    # [code: main.py#L149]
    valid = (d_star > 0) & (d_star < D_max)
    total = 0
    for prediction, weight in zip(disparities, weights):
        # Smooth L1 averaged over valid pixels  [eq. 6]
        # [code: models/loss.py#L4-L9]
        diff = prediction[valid] - d_star[valid]
        total = total + weight * smooth_l1_mean(diff)
    return total

Hyperparameters

Results

Author-reported. Setup: Scene Flow Finalpass, with the ablations trained by the authors at their batch size §4.1, tab. 2. KITTI test numbers come from the evaluation server §4.6. Timing is for 480×640 inputs on one Nvidia TITAN Xp tab. 2.

Scene Flow ablation tab. 2. "Base" removes the stacked hourglasses §4.3.

KITTI validation tab. 3.

KITTI 2015 test tab. 4. D1 outlier percentages, with time in seconds.

KITTI 2012 test tab. 5. Noc and All variants, with time in seconds.

Other results, in words. The group count sweep §4.3, tab. 2 improves up to about 40 groups, after which the EPE stays near 1.2 px. Figure 3 shows the gap between the group-wise model (Gwc-Cat) and the concatenation-only model (Cat) widening as the 3D network's channel count shrinks §4.5, fig. 3. The visualisations of depth, prediction and error maps are in fig. 4. The paper does not give a numeric summary for them.

Paper vs code

Code is at commit 0987074f48b098c6e975a58ca5aacf233b0fb0eb. It was found via the link in the abstract, and its owner is the first author.

• Scene Flow test crop. The paper feeds the full 960×540 images §4.2. The code crops the bottom-right 960×512 region, dropping the top 28 rows of a 540-row image sceneflow_dataset.py. The reported numbers therefore come from a slightly smaller test input than the text says.
• ReLU after the second pre-hourglass convolution. The paper says the pre-hourglass module has four 3D convolutions with BN and ReLU §3.1, and the footnote of tab. 1 says every 3D convolution has ReLU unless marked otherwise. The code's second block has no ReLU after its second convolution, and its residual sum is not followed by one gwcnet.py gwcnet.py.
• First pre-hourglass convolution. Tab. 1 lists it as 32×32, but its input is the 64-channel volume. The code maps 64 to 32 channels gwcnet.py.
• Loss mask. Eq. 5 has no validity mask. The code trains only on pixels with 0 < d* < D_max main.py and averages the smooth L1 loss over those pixels loss.py. Paper and code differ here.
• Training crops. Both datasets are randomly cropped to 512×256 during training sceneflow_dataset.py kitti_dataset.py. The paper does not state this.
• Dilation. The paper says "half dilation settings" §3.1. The code dilates only the last residual stage, with dilation 2 [code: models/gwcnet.py#L27]. This is the code's reading of the phrase.
• Initialisation and weight decay. Not in the paper. The code uses He-style normal initialisation gwcnet.py and no weight decay main.py.
• Number of groups. Fixed at 40 in the model gwcnet.py and not exposed as a flag, so changing it means editing the code.
• Inference outputs. The paper says output modules 0 to 2 can be dropped at inference §3.3. The code does this gwcnet.py. No discrepancy.
• Learning-rate schedule. Matches the text; the code halves cumulatively experiment.py. No discrepancy.

No discrepancies were found for the group-wise correlation formula, the soft argmin, the loss weights, the smooth L1 definition, the concatenation channel counts, or the D1 outlier rule.

Pitfalls

• Mean, not sum. Eq. 3 divides by the channels per group (8 here), so the correlation scale depends on group size, not on the number of groups §3.2, eq. 3 submodule.py. Use the mean.
• Zeros outside the image. For disparities with x < d, the group-wise and concatenation volumes stay at zero, and nothing masks them out before aggregation submodule.py submodule.py. The paper does not describe this case.
• Upsampling is unspecified. The paper says "upsampled" §3.4. The code uses trilinear interpolation and does not set an alignment option gwcnet.py. The paper states neither the mode nor the alignment.
• Table 4 key. The PSMNet row of tab. 4 cites kendall2017end, the key the text uses for GC-Net §2.2, tab. 4. The paper does not say which reference the row is meant to carry.
• Loss scale. The paper does not say whether smooth L1 is summed or averaged over pixels. The code averages over the valid pixels per output loss.py. The result depends on this choice.
• Wording vs table. §4.4 says the new hourglass "increases EPE by 7.8%" on Scene Flow. Tab. 2 shows EPE falling from 0.876 to 0.808, which is a reduction, so the table is the reliable source §4.4, tab. 2.
• Retrained baseline. The PSMNet row in tab. 2 is the authors' retrain with their batch size and evaluation settings, not the published PSMNet result tab. 2 caption.
• Validation noise. The authors say the KITTI validation results are unstable, so they fine-tuned three times and kept the best model §4.6. The selection procedure is not in the code.
• Hidden costs. The volume has N_g channels over D_max/4 disparity slices at 1/4 resolution §3.2. For D_max = 192 that is 48 slices. The paper gives timings only for 480×640 inputs on one TITAN Xp tab. 2, and memory use is not reported.
• Gwc1 is full correlation. Gwc1-Base (1 group × 320 channels) is full correlation in the paper's own terms §3.2, so it is the reference point of the group-count sweep tab. 2.

Lineage

• Builds on: PSMNet [chang2018pyramid], whose feature extractor and 3D stacked hourglass the paper modifies. The paper does not give an arXiv id, and the PSMNet card is not in this wiki §3.1, §3.3. GC-Net [kendall2017end] for the soft argmin and the concatenation volume §2.2, §3.4. DispNetC [mayer2016large] for full correlation §2.2.
• Compared against: DispNetC, GC-Net, CRL, iResNet, PSMNet, SegStereo and MC-CNN-acrt tab. 4, tab. 5.
• Superseded by: not stated in this wiki.