双目深度算法——基于Transformer的方法（STTR）

双目深度算法——基于Transformer的方法（STTR）

STTR是STereo TranformeR的缩写，原论文名为《Revisiting Stereo Depth Estimation From a Sequence-to-Sequence Perspectivewith Transformers》，发表于2021年，据我了解这应该是第一篇使用Transformer进行双目视差估计的方法，打破了基于Correlation或者Cost Volume进行视差估计的方法论，在论文的摘要中，作者提到该方法主要有三大优势：（1）解放了视差的限制；（2）明确定义了遮挡区域；（3）保证了匹配的唯一性，这篇文章实验做得非常充分，开源代码也些得很好，下面我们结合代码和文中的实验来详细学习下这篇论文。

1. 网络构架

网络整体结果如下图所示，主要包括三部分：Feature Extractor、Transformer和Context Adjustment Layer，其中Feature Extractor主要用于特征提取，Transformer通过Attention计算视差，Context Adjustment Layer用于后处理。

在STTR的代码中也对应地将网络抽象为几部分：

def forward(self, x: NestedTensor):
    """
    :param x: input data
    :return:
        a dictionary object with keys
        - "disp_pred" [N,H,W]: predicted disparity
        - "occ_pred" [N,H,W]: predicted occlusion mask
        - "disp_pred_low_res" [N,H//s,W//s]: predicted low res (raw) disparity
    """
    bs, _, h, w = x.left.size()

    # extract features
    feat = self.backbone(x)  # concatenate left and right along the dim=0
    tokens = self.tokenizer(feat)  # 2NxCxHxW
    pos_enc = self.pos_encoder(x)  # NxCxHx2W-1

    # separate left and right
    feat_left = tokens[:bs]
    feat_right = tokens[bs:]  # NxCxHxW

    # downsample
    if x.sampled_cols is not None:
        feat_left = batched_index_select(feat_left, 3, x.sampled_cols)
        feat_right = batched_index_select(feat_right, 3, x.sampled_cols)
    if x.sampled_rows is not None:
        feat_left = batched_index_select(feat_left, 2, x.sampled_rows)
        feat_right = batched_index_select(feat_right, 2, x.sampled_rows)

    # transformer
    attn_weight = self.transformer(feat_left, feat_right, pos_enc)

    # regress disparity and occlusion
    output = self.regression_head(attn_weight, x)

    return output

其中backbone为encoder部分，tokenizer为decoder部分。

1.1 Feature Extractor

Feature Extractor主要分为Encoder和Decoder两部分，其中Encoder部分使用的是类似Hourglass的结构，在Decoder部分使用的是转置卷积和Dense Block，特征提取的网络结构就不在此展开，其主要作用就是从原始的图像输入中提取图像特征，特征图大小和原始图像大小相同，但是每个像素变成了一个长为$C_e$的特征向量。

尽管论文中是讲，基于Transformer的网络结构没有视差的限制，但是由于特征提取使用的CNN网络，因此计算Self Attention和Cross Attention使用的特征向量还是从图像的一个局部区域（感受野）抽象出来的。

1.2 Transformer

Transformer部分结构如下图所示：

这种反复叠加Self-Attention和Cross-Attention的机制是参考的特征匹配算法SuperGlue，对该算法不熟悉的同学可以参考视觉SLAM总结——SuperPoint / SuperGlue，在该算法的论文中提到，使用这种反复叠加的机制的目的是模仿人类在完成此类任务时的行为，我们寻找两幅图像上相似的像素点时也是先在其中一幅图像上对比该像素与其他像素的区别，然后再在另一幅图像上去尝试寻找最接近的像素。

本论文使用的是带残差的多头注意力机制，公式如下：$${\mathcal{Q}}_h=W_{{\mathcal{Q}}_h}{e}_I+b_{{\mathcal{Q}}_h}$$$${\mathcal{K}}_h=W_{{\mathcal{K}}_h}{e}_I+b_{{\mathcal{K}}_h}$$$${\mathcal{V}}_h=W_{{\mathcal{V}}_h}{e}_I+b_{{\mathcal{V}}_h}$$$${\alpha }_h=\mathrm{softmax}\left(\frac{{\mathcal{Q}}_h^T{K}_h}{\sqrt{C_h}}\right)$$$$V_{\mathcal{O}}=W_{\mathcal{O}}\text{ Concat }\left({\alpha }_1{\mathcal{V}}_1,\dots ,{\alpha }_{N_h}{\mathcal{V}}_{N_h}\right)+b_{\mathcal{O}}$$$$e_I=e_I+{\mathcal{V}}_{\mathcal{O}}$$其中$W_{{\mathcal{Q}}_h},W_{{\mathcal{K}}_h},W_{{\mathcal{V}}_h}\in {\mathbb{R}}^{C_h\times C_h},b_{{\mathcal{Q}}_h},b_{{\mathcal{K}}_h},b_{{\mathcal{V}}_h}\in {\mathbb{R}}^{C_h}$以及$W_{\mathcal{O}}\in {\mathbb{R}}^{C_e\times C_e},b_{\mathcal{O}}\in {\mathbb{R}}^{C_e}$，这就是普通的Attention计算公式，我们就不在此赘述，细节不清楚的同学可以参考计算机视觉算法——Transformer学习笔记，作者在实现Attention机制时是因为加入了相对位置编码和注意力掩膜，所以是继承原始pytorch中的MultiheadAttention类重新实现了下，但是这一部分基本的操作是保持不变的：

# project to get qkv
if torch.equal(query, key) and torch.equal(key, value):
    # self-attention
    q, k, v = F.linear(query, self.in_proj_weight, self.in_proj_bias).chunk(3, dim=-1)
else
	...
	
# reshape	
q = q.contiguous().view(w, bsz, self.num_heads, head_dim)  # WxNxExC
if k is not None:
    k = k.contiguous().view(-1, bsz, self.num_heads, head_dim)
if v is not None:
    v = v.contiguous().view(-1, bsz, self.num_heads, head_dim)

# compute attn weight
attn = torch.einsum('wnec,vnec->newv', q, k)  # NxExWxW'
attn = F.softmax(attn, dim=-1)
v_o = torch.bmm(attn.view(bsz * self.num_heads, w, w), v.permute(1, 2, 0, 3).view(bsz * self.num_heads, w, head_dim))  # NxExWxW', W'xNxExC -> NExWxC
v_o = v_o.reshape(bsz, self.num_heads, w, head_dim).permute(2, 0, 1, 3).reshape(w, bsz, embed_dim)
v_o = F.linear(v_o, self.out_proj.weight, self.out_proj.bias)

# average attention weights over heads
attn = attn.sum(dim=1) / self.num_heads

作者在补充材料中可视化了不同层的注意力分布结果，如下图所示：

可以看到随着层数的加深，Self-Attention和Cross-Attention都逐渐收敛到正确的视差位置

1.2.1 Relative Positional Encoding

为了保证算法在大范围无纹理区域也能够有合理的视差估计，作者提到需要在输入中加入与数据无关位置信息，即$$e=e_I+e_p$$上式展开后可以获得$$\begin{array}{c}{\alpha }_{i,j}=\underset{(1)\text{ data-data }}{\underbrace{e_{I,i}^T{W}_{\mathcal{Q}}^T{W}_{\mathcal{K}}{e}_{I,j}}}+\underset{(2)\text{ data-position }}{\underbrace{e_{I,i}^T{W}_{\mathcal{Q}}^T{W}_K{e}_{p,j}}}+\\ \underset{\text{(3) position-data }}{\underbrace{e_{p,i}^T{W}_{\mathcal{Q}}^T{W}_{\mathcal{K}}{e}_{I,j}}}+\underset{\text{(4) position-position }}{\underbrace{e_{p,i}^T{W}_{\mathcal{Q}}^T{W}_{\mathcal{K}}{e}_{p,j}}}.\end{array}$$我们注意到上式中第四项仅仅取决于像素点位置，而和图像信息完全无关了，是不必存在的一项，我们将第四项移除后得到最后的结果：$$\begin{array}{r}{\alpha }_{i,j}=\underset{(1)\text{ data-data }}{\underbrace{e_{I,i}^T{W}_{\mathcal{Q}}^T{W}_{\mathcal{K}}{e}_{I,j}}}+\\ \underset{\text{(2) data-position }}{\underbrace{e_{I,i}^T{W}_{\mathcal{Q}}^T{W}_K{e}_{p,i-j}}}+\underset{\text{(3) position-data }}{\underbrace{e_{p,i-j}^T{W}_{\mathcal{Q}}^T{W}_{\mathcal{K}}{e}_{I,j}}},\end{array}$$在具体实现时，就是在计算完${\alpha }_h$后加上Positional Encoding

# 0.3 s
attn_feat_pos = torch.einsum('wnec,wvec->newv', q, k_r)  # NxExWxW'
attn_pos_feat = torch.einsum('vnec,wvec->newv', k, q_r)  # NxExWxW'

# 0.1 s
attn = attn_feat + attn_feat_pos + attn_pos_feat

在论文的Ablation实验中，作者对比了有Positional Encoding和没有Positional Encoding的特征图的区别：

可以看到，加入Positional Encoding的特征图，在若纹理区域仍然保持着和边界相关的纹理，这对于最后的视差匹配是由帮助的。在精度对比的实验中，也证明了加入Positional Encoding是有效的。

但在这里我有疑惑的一点的是，在ViT论文的实验中，发现Absolute Positional Encoding和Relative Positional Encoding最后的结果是差别不大，在本文论中，也并没有对比作者提出的Relative Positional Encoding相对简单的Absolute Positional Encoding的区别，这一点值得考究。

1.2.2 Optimal Transport

Optimal Transport算法是一种可微分的二分图匹配算法，这一部分的内容作者是完全参考SuperGlue的做法，包括Dustbin的设置（本文在匹配到Dustbin的数据上加了一个可学习参数$\varphi$），这里感兴趣的同学可以参考视觉SLAM总结——SuperPoint / SuperGlue

1.2.3 Attention Mask

Attention Mask只在最后一层Cross Attention中用到了，其基本原理如下图所示：

图（a）中左图中的桌角在右图中绝对不可能出现在右图虚线的左边，因此我们在进行Cross Attention的计算过程中，可以减少一些计算量，如图（b）中这样一个Attention Mask，黑色区域的Cross Attention就不再需要计算，具体实现是通过下面函数生成mask

@torch.no_grad()
def _generate_square_subsequent_mask(self, sz: int):
    """
    Generate a mask which is upper triangular
    :param sz: square matrix size
    :return: diagonal binary mask [sz,sz]
    """
    mask = torch.triu(torch.ones(sz, sz), diagonal=1)
    mask[mask == 1] = float('-inf')
    return mask

然后在最后一层Cross Attention上将生成的mask加到原始的Attention结果上

attn_mask = attn_mask[None, None, ...]
attn += attn_mask

其中就是在不需要计算的区域加上一个负无穷的值，这样在Softmax后这些区域的权重就自然变成了0，在论文的Ablation实验中，这个模块也可以带来精度的提高。

1.2.4 Raw Disparity and Occlusion Regression

在基于Correlation和Cost Volume的方法中，求视差通常采用的是加权平均，而本文由于采用的是匹配的方法，因此最终的结果是通过Modified Winner-Take-All的方式获得的，如何Modified的呢？作者先从分配矩阵$\mathcal{T}$中取得一个$3\times 3$的Patch进行归一化$${\tilde{t}}_l=\frac{t_l}{{\sum }_{l\in {\mathcal{N}}_3(k)}{t}_l},\text{ for }l\in {\mathcal{N}}_3(k)$$然后使用归一化的结果对视差进行加权平均$${\tilde{d}}_{raw}(k)=\sum _{l\in {\mathcal{N}}_3(k)}{d}_l{\tilde{t}}_l$$这个$3\times 3$的Patch也正好代表着网络对于当前匹配结果的一个确定程度，因此对其求反就可以作为一个被遮挡的概率$$p_{occ}(k)=1-\sum _{l\in {\mathcal{N}}_3(k)}{t}_l.$$但是这里求得的视差和遮挡概率都是最原始的结果，后面还需要通过Context Adjustment Layer进行Refine。

1.3 Context Adjustment Layer

Context Adjustment Layer的主要目的是对输出的视差和遮挡不确定度进行Refine，其结构如下图所示：

大致原理就是将原始图像作为数据，利用原始图像的信息通过卷积对结果进行修正，这种方法在基于Correlation和Cost Volume的方法中都有用到，文中有对比使用CAL和不使用CAL的区别：

2. 损失函数

STTR算法中损失函数构建得相对复杂，一共由四部分组成，如下：$$L=w_1{L}_{rr}+w_2{L}_{d1,r}+w_3{L}_{d1,f}+w_4{L}_{be,f}$$其中
第一部分$L_{rr}$是基于分配矩阵$\mathcal{T}$构建的损失：$$t_{\mathrm{i}}^{\ast }=\mathrm{interp}\left({\mathcal{T}}_{\mathrm{i}},p_{\mathrm{i}}-d_{gt,i}\right)$$$$L_{rr}=\frac{1}{N_{\mathcal{M}}}\sum _{i\in \mathcal{M}}-\log \left(t_i^{\ast }\right)+\frac{1}{N_{\mathcal{U}}}\sum _{i\in \mathcal{U}}-\log \left(t_{i,\varphi }\right)$$其中$t_i^{\ast }$为插值后的分配矩阵$\mathcal{T}$的值，$t_{i,\varphi }$为分配到未匹配区域的值。

第二部分$L_{d1,r}$和$L_{d1,f}$分别为原始视差和Refine后的视差的L1损失。

第三部分$L_{\mathrm{be},f}$为判断是否为遮挡区域的二分类交叉熵损失。

3. 实验及实验结果

为了减小内存占用，作者采用的是Gradient Checkpointing技术，该技术是使得一些临时变量在前向推导时不进行存储，而时在反向导数传递时重新计算，这样就是用时间换空间，在训练大型模型时是一种常见的处理手段，具体实现其实就是使用了torch.utils.checkpoint这个库，如下所示：

def _alternating_attn(self, feat: torch.Tensor, pos_enc: torch.Tensor, pos_indexes: Tensor, hn: int):
    """
    Alternate self and cross attention with gradient checkpointing to save memory
    :param feat: image feature concatenated from left and right, [W,2HN,C]
    :param pos_enc: positional encoding, [W,HN,C]
    :param pos_indexes: indexes to slice positional encoding, [W,HN,C]
    :param hn: size of HN
    :return: attention weight [N,H,W,W]
    """

    global layer_idx
    # alternating
    for idx, (self_attn, cross_attn) in enumerate(zip(self.self_attn_layers, self.cross_attn_layers)):
        layer_idx = idx

        # checkpoint self attn
        def create_custom_self_attn(module):
            def custom_self_attn(*inputs):
                return module(*inputs)

            return custom_self_attn

        feat = checkpoint(create_custom_self_attn(self_attn), feat, pos_enc, pos_indexes)

        # add a flag for last layer of cross attention
        if idx == self.num_attn_layers - 1:
            # checkpoint cross attn
            def create_custom_cross_attn(module):
                def custom_cross_attn(*inputs):
                    return module(*inputs, True)

                return custom_cross_attn
        else:
            # checkpoint cross attn
            def create_custom_cross_attn(module):
                def custom_cross_attn(*inputs):
                    return module(*inputs, False)

                return custom_cross_attn

        feat, attn_weight = checkpoint(create_custom_cross_attn(cross_attn), feat[:, :hn], feat[:, hn:], pos_enc,
                                       pos_indexes)

    layer_idx = 0
    return attn_weight

在算法对比部分，作者显示对比了模型的泛化性，作者将所有模型在Sense Flow数据集上训练，然后直接在KITTI数据集上推理，结果如下图所示：

可见STTR取得了较好的模型泛化性，但是当作者将算法在KITTI数据集上进行Fine Tune后对比结果就没那么理想了：

作者解释到，KITTI数据集训练集实在太小且视差都比较小，的确由于Transfomer的模型通常需要较大的数据量进行训练，我也觉得当数据量起来后，基于Transformer的模型一定会有更加明显的优势，最后作者还进行了一系列剪枝的实验，但是即使轻量化后STTR也很难达到事实的效果，我相信后面肯定还会有工作对此进行优化。