1 概述

本文主要結合代碼對Selective的創新點進行針對性講解，相關的背景知識可以參考我寫的另兩篇文章論文閱讀-RaftStereo和論文閱讀-IGEV。
SelectiveStereo的創新點總結來說就只有一項，在RaftStereo和IGEV的基礎上，提出了分別提取圖像高頻信息和低頻信息并融合的迭代算子SRU(Selective Recurrent Unit)，另一個Contextual Spatial Attention(CSA)模塊是為其服務的。
SelectiveStereo對RaftStereo和IGEV都進行了改進，下面以IGEV為例進行具體說明。

2 模塊

2.1 SelectiveIGEV和IGEV的差異

從代碼上看，SelectiveStereo和IGEV的差異只有兩處。一處是在計算inp_list的時候，SelectiveStereo用到了通道注意力self.cam，并使用空間注意力self.sam進一步得到了注意力圖att；另一處是在使用self.update_block進行計算的時候，SelectiveIGEV額外多了att作為輸入。
SelectiveStereo的相關代碼為

class IGEVStereo(nn.Module):def forward(self, image1, image2, iters=12, flow_init=None, test_mode=False):# 此處省略部分代碼...cnet_list = self.cnet(image1, num_layers=self.args.n_gru_layers)net_list = [torch.tanh(x[0]) for x in cnet_list]inp_list = [torch.relu(x[1]) for x in cnet_list]# 不同點1inp_list = [self.cam(x) * x for x in inp_list]att = [self.sam(x) for x in inp_list]# 此處省略部分代碼...for itr in range(iters):disp = disp.detach()geo_feat = geo_fn(disp, coords)with autocast(enabled=self.args.mixed_precision, dtype=getattr(torch, self.args.precision_dtype, torch.float16)):# 不同點2net_list, mask_feat_4, delta_disp = self.update_block(net_list, inp_list, geo_feat, disp, att)

IGEV的相關代碼為

class IGEVStereo(nn.Module):def forward(self, image1, image2, iters=12, flow_init=None, test_mode=False):# 此處省略部分代碼...cnet_list = self.cnet(image1, num_layers=self.args.n_gru_layers)net_list = [torch.tanh(x[0]) for x in cnet_list]inp_list = [torch.relu(x[1]) for x in cnet_list]# 不同點1inp_list = [list(conv(i).split(split_size=conv.out_channels//3, dim=1)) for i,conv in zip(inp_list, self.context_zqr_convs)]# 此處省略部分代碼...for itr in range(iters):disp = disp.detach()geo_feat = geo_fn(disp, coords)with autocast(enabled=self.args.mixed_precision, dtype=getattr(torch, self.args.precision_dtype, torch.float16)):# 不同點2net_list, mask_feat_4, delta_disp = self.update_block(net_list, inp_list, geo_feat, disp, iter16=self.args.n_gru_layers==3, iter08=self.args.n_gru_layers>=2)

2.2 上下文空間注意力

為實現不同感受野和頻率信息的融合，上下文空間注意力（Contextual Spatial Attention，簡稱CSA）模塊通過提取上下文信息作為引導的多層級注意力圖譜。如圖2-1所示，CSA可劃分為兩個子模塊：通道注意力增強（Channel Attention Enhancement，簡稱CAE）和空間注意力提取器（Spatial Attention Extractor，簡稱SAE）。這兩個子模塊源自CBAM框架，作者進行了簡化以更好地適配立體匹配任務。

圖2-1 上下文注意力示意圖

2.2.1 通道注意力

通道注意力的代碼實現如下所示，結合圖2-1左側的示意圖應該比較容易理解。在給定上下文信息圖 $c\in R^{C\times H\times W}$ 的情況下，首先對空間維度進行平均池化和最大池化操作，分別得到兩個特征圖$f_{avg}$和$f_{max}\in R^{C\times 1\times 1}$。接著通過兩個獨立的卷積層對這些特征圖進行特征變換，也就是代碼中的self.fc，是一個bottleneck結構，如果改成殘差的形式就和yolo中的bottleneck一致了。隨后將這兩個特征圖相加，并使用sigmoid函數將其轉換為0到1之間的權重$M_c\in R^{C\times 1\times 1}$。最后通過逐元素乘積運算，初始特征圖能夠捕捉哪些通道具有數值高的特征值需要增強，哪些通道具有數值低的特征值需要抑制。

class ChannelAttentionEnhancement(nn.Module):def __init__(self, in_planes, ratio=16):super(ChannelAttentionEnhancement, self).__init__()self.avg_pool = nn.AdaptiveAvgPool2d(1)self.max_pool = nn.AdaptiveMaxPool2d(1)self.fc = nn.Sequential(nn.Conv2d(in_planes, in_planes // 16, 1, bias=False),nn.ReLU(),nn.Conv2d(in_planes // 16, in_planes, 1, bias=False))self.sigmoid = nn.Sigmoid()def forward(self, x):avg_out = self.fc(self.avg_pool(x))max_out = self.fc(self.max_pool(x))out = avg_out + max_outreturn self.sigmoid(out)

回看2.1節中SelectiveStereo對應的代碼，self.cam生成了inp_list，而IGEV中是通過簡單的卷積生成的inp_list，因此SelectiveStereo的inp_list特征表達能力會更強一些。
那么這個inp_list的作用是什么呢？看過我的另一篇博客論文閱讀-IGEV的話可以知道inp_list對應于ConvGRU中的偏置項，而在SelectiveStereo中，作者選擇使用了RaftConvGRU，那么inp_list就不是偏置項了，直接作為卷積的輸入了。我覺得作者應該是認為經過sigmoid的結果作為偏置項不太合適才選擇了RaftConvGRU。
這里也把ConvGRU和RaftConvGRU的代碼放出來，方便讀者查看。
RaftConvGRU代碼如下，inp_list被concat在x當中。

class RaftConvGRU(nn.Module):def __init__(self, hidden_dim=128, input_dim=256, kernel_size=3, dilation=1):super(RaftConvGRU, self).__init__()self.convz = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=(kernel_size+(kernel_size-1)*(dilation-1))//2, dilation=dilation)self.convr = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=(kernel_size+(kernel_size-1)*(dilation-1))//2, dilation=dilation)self.convq = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=(kernel_size+(kernel_size-1)*(dilation-1))//2, dilation=dilation)def forward(self, h, x):hx = torch.cat([h, x], dim=1)z = torch.sigmoid(self.convz(hx))r = torch.sigmoid(self.convr(hx))q = torch.tanh(self.convq(torch.cat([r*h, x], dim=1)))h = (1-z) * h + z * qreturn h

ConvGRU代碼如下，inp_list就是[cz, cr, cq]。

class ConvGRU(nn.Module):def __init__(self, hidden_dim, input_dim, kernel_size=3):super(ConvGRU, self).__init__()self.convz = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size//2)self.convr = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size//2)self.convq = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size//2)def forward(self, h, cz, cr, cq, *x_list):x = torch.cat(x_list, dim=1)hx = torch.cat([h, x], dim=1)z = torch.sigmoid(self.convz(hx) + cz)r = torch.sigmoid(self.convr(hx) + cr)q = torch.tanh(self.convq(torch.cat([r*h, x], dim=1)) + cq)h = (1-z) * h + z * qreturn h

2.2.2 空間注意力

空間注意力的代碼如下所示，結合圖2-1右側的示意圖應該比較容易理解。采用和通道注意力相同的池化操作，但這次將池化維度調整為通道維度，對應于代碼中的avg_out和max_out。隨后將這些池化結果拼接成$R^{2\times H\times W}$的圖像，并通過帶有sigmoid函數的卷積層生成最終注意力圖。回顧之前的處理流程可以發現，該注意力圖在需要高頻信息的區域具有高權重值，因為這些信息在上下文信息中具有較高的特征價值。而對需要低頻信息的區域則賦予較低權重。總體而言，該注意力圖能明確區分不同頻率信息需求的區域。這個圖有點像邊緣圖，物體的邊緣即使圖像中的高頻信息區域。

class SpatialAttentionExtractor(nn.Module):def __init__(self, kernel_size=7):super(SpatialAttentionExtractor, self).__init__()self.samconv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2, bias=False)self.sigmoid = nn.Sigmoid()def forward(self, x):avg_out = torch.mean(x, dim=1, keepdim=True)max_out, _ = torch.max(x, dim=1, keepdim=True)x = torch.cat([avg_out, max_out], dim=1)x = self.samconv(x)return self.sigmoid(x)

2.3 SRU

SRU看代碼就非常容易理解了，它設計了兩個GRU。一個GRU的kernel_size小，捕獲圖像中的高頻信息，只需要小的感受野，如邊緣，最終的結果和表示圖像中高頻區域的注意力att相乘；另一個GRU的kernel_size大，捕獲圖像中的低頻信息，如平面，最終的結果和表示圖像中低頻區域的注意力1-att相乘。如果既不是高頻也不是低頻的，att就會趨于0.5，那么就會結合self.small_gru和self.large_gru兩者的信息。

class SelectiveConvGRU(nn.Module):def __init__(self, hidden_dim=128, input_dim=256, small_kernel_size=1, large_kernel_size=3):super(SelectiveConvGRU, self).__init__()self.small_gru = RaftConvGRU(hidden_dim, input_dim, small_kernel_size)self.large_gru = RaftConvGRU(hidden_dim, input_dim, large_kernel_size)def forward(self, att, h, *x):x = torch.cat(x, dim=1)h = self.small_gru(h, x) * att + self.large_gru(h, x) * (1 - att)return h

SRU的示意圖如下圖2-2所示。

圖2-2 SRU示意圖

3 效果

加入SelectiveStereo之后，RaftStereo和IGEV的效果在不同的update次數下都有提升，如下表3-1所示。

表3-1 不同update次數下SelectiveStereo的效果

IGEV和SelectiveIGEV可視化效果對比如下圖3-1所示，SelectiveIGEV的細節部分效果明顯更好。

圖3-1 IGEV和SelectiveIGEV可視化效果對比

SelectiveStereo在SceneFlow數據集和不同模型的EPE指標對比如下表3-2所示。

表3-2 SelectiveStereo在SceneFlow數據集和不同模型的EPE指標對比

SelectiveStereo在KITTI數據集和不同模型的EPE指標對比如下表3-3所示。

表3-3 SelectiveStereo在KITTI數據集和不同模型的EPE指標對比

整體來說，使用了SelectiveStereo后效果方面會得到全方面的提升，不過速度方面是會下降的，因為引入了把一個GRU標成了大小兩個GRU，且多了注意力的計算步驟。

參考資料

[1] Selective-Stereo: Adaptive Frequency Information Selection for Stereo Matching
[2] https://github.com/Windsrain/Selective-Stereo.git
[3] https://github.com/gangweix/IGEV.git