GMFlow 原理+代码解析

GMFlow: Learning Optical Flow via Global Matching（CVPR 2022）

GMFlow 原理

Motivation

• 回归光流：过去几年中具有代表性的光流学习框架的核心估计方式没有太大本质区别，即利用卷积Conv从图像的局部相关性中回归光流值（回归问题）。如基于深度学习估计光流的开创性工作 FlowNet，设计了一个卷积神经网络架构，可以直接将两张视频帧作为输入，并输出稠密的光流。基于学习方法的进一步发展使得光流估计的效果在稳固提升，同时很多不同的光流网络结构被提出。然而，如果我们仔细考量下各种不同的光流网络，可以发现，基本的光流估计方式与最早的 FlowNet 无太大本质区别，都是利用卷积从局部相关性中回归光流。
  

• 迭代回归光流：因为回归光流的方式由于其内在的局部性，难以处理光流领域中长期存在的一个挑战：大运动的估计。为了缓解这一问题，当前的代表性框架 RAFT 利用大量的迭代改价来逐步提升光流的预测效果。相较之前的算法，RAFT 取得了巨大的性能提升，在最近一两年中各种 RAFT 变体层出不穷。然而，尽管这种迭代框架的取得了出色的性能，但由于本质上是一种序列化的处理方式，它也带来了线性的推理时间增长，使其难以做速度上的优化。
  

• 全局匹配光流：GMFlow主要要验证这样一个假设：无需大量迭代，同样可以取得很好的光流估计效果，同时速度更快。
  

假设我们要在第二张图中找第一张图蓝色点的对应点，一般地，我们通常会浏览第二张图中的所有像素点，并比较这些点与蓝色点的相似度，最终将相似度最高的点作为对应点，即黄色点。这一观察启发我们重新审视光流这一任务的本质：光流究竟是一个回归问题还是匹配问题？

光流估计直观上是一个匹配问题，目的是寻找对应的像素。 为了实现这一点，可以比较每个像素的特征相似性，并确定具有最高相似性的对应像素。这样的过程要求特征具有足够的辨别性。将图像本身的空间上下文和来自另一张图像的信息聚合在一起，可以直观地缓解歧义，提高其辨别性。这样的设计理念使得稀疏特征匹配框架取得了巨大的成就。我们还从另外一个相关任务中得到一些启发，即两张图像 (未必是视频帧) 之间的稀疏对应关系，往往用于运动恢复结构和相机位姿估计等应用。这个任务的特点是通常两张图片之间的视角差异较大。我们注意到在主流的稀疏框架中，对应点往往是通过匹配得到的。

基于这些观察，GMFlow提出将光流重新定义为一个全局匹配问题，以期能更好地解决大运动这一难题。

Method

Core Idea：在本文中，将Flow Estimation重新定义为一个全局匹配问题，即通过直接比较所有特征之间的相似度来得到稠密对应关系。这种问题定义依赖于较强的特征，为此利用了Transformer来实现。

给定两个连续的视频帧I1和I2：

• 首先采用权重共享的Conv Layer提取 8 倍下采样的稠密特征F1,F2$\in R^{H\times W\times D}$。

• 然后送入Transformer进一步提取特征。

• 接着通过全局相关性计算所有点与点之间的相关性，通过矩阵乘法现，其中相关矩阵C中的每个元素表示F1中$pixe{l}_1=(i,j)$和F2中$pixe{l}_2=(k,l)$的相关性值，1/D为归一化因子，避免点积运算后值较大：
  $$C=\frac{F_1{F}_2^T}{\sqrt{D}}\in R^{H\times W\times H\times W}$$

• 再使用 Softmax 匹配层获得光流Flow。想要确定两个feature中pixel的对应关系，一种可行的方法是直接取C矩阵中相关性值最高的位置。然而，这个操作是不可微的，阻碍了端到端训练。为了解决这个问题，我们使用了一个可微匹配层，即经过Softmax层归一化C矩阵的最后两个维度，这给了我们一个匹配的概率分布M，表示F1中的每个位置相对于F2中的所有位置是否匹配：
  $$M=Softmax(C)\in R^{H\times W\times H\times W}$$

• 然后，将像素网格$G\in R^{H\times W\times 2}$的二维坐标加权平均，得到对应的G与匹配的分布：$\hat{G}=MG\in R^{H\times W\times 2}$。最后，通过计算对应像素坐标的差值就可以得到光流V：$V=\hat{G}-G\in R^{H\times W\times 2}$。

• 最后引入了一个额外的Self-Attention Layer，通过考虑图像特征自相似性，将匹配像素中的高质量FLow预测传播到不匹配的像素，来缓解遮挡的和边界外的像素点的问题：
  

GMFlow 代码

核心代码分析

初始化：CNNEncoder、FeatureTransformer、FeatureFlowAttention、CNNUpsampler 四部分组成。

class GMFlow(nn.Module):
    def __init__(self,
                 num_scales=1,
                 upsample_factor=8,
                 feature_channels=128,
                 attention_type='swin',
                 num_transformer_layers=6,
                 ffn_dim_expansion=4,
                 num_head=1,
                 **kwargs,
                 ):
        super(GMFlow, self).__init__()

        self.num_scales = num_scales
        self.feature_channels = feature_channels
        self.upsample_factor = upsample_factor
        self.attention_type = attention_type
        self.num_transformer_layers = num_transformer_layers

        # CNN backbone
        self.backbone = CNNEncoder(output_dim=feature_channels, num_output_scales=num_scales)

        # Transformer
        self.transformer = FeatureTransformer(num_layers=num_transformer_layers,
                                              d_model=feature_channels,
                                              nhead=num_head,
                                              attention_type=attention_type,
                                              ffn_dim_expansion=ffn_dim_expansion,
                                              )

        # flow propagation with self-attn
        self.feature_flow_attn = FeatureFlowAttention(in_channels=feature_channels)

        # convex upsampling: concat feature0 and flow as input
        self.upsampler = nn.Sequential(nn.Conv2d(2 + feature_channels, 256, 3, 1, 1),
                                       nn.ReLU(inplace=True),
                                       nn.Conv2d(256, upsample_factor ** 2 * 9, 1, 1, 0))

forward推理：

def forward(self, img0, img1,
                attn_splits_list=None,
                corr_radius_list=None,
                prop_radius_list=None,
                pred_bidir_flow=False,
                **kwargs,
                ):

        results_dict = {}
        flow_preds = []

        img0, img1 = normalize_img(img0, img1)  # [B, 3, H, W]

        # resolution low to high
        feature0_list, feature1_list = self.extract_feature(img0, img1)  # list of features

        flow = None

        assert len(attn_splits_list) == len(corr_radius_list) == len(prop_radius_list) == self.num_scales

        for scale_idx in range(self.num_scales):
            feature0, feature1 = feature0_list[scale_idx], feature1_list[scale_idx]

            if pred_bidir_flow and scale_idx > 0:
                # predicting bidirectional flow with refinement
                feature0, feature1 = torch.cat((feature0, feature1), dim=0), torch.cat((feature1, feature0), dim=0)

            upsample_factor = self.upsample_factor * (2 ** (self.num_scales - 1 - scale_idx))

            if scale_idx > 0:
                flow = F.interpolate(flow, scale_factor=2, mode='bilinear', align_corners=True) * 2

            if flow is not None:
                flow = flow.detach()
                feature1 = flow_warp(feature1, flow)  # [B, C, H, W]

            attn_splits = attn_splits_list[scale_idx]
            corr_radius = corr_radius_list[scale_idx]
            prop_radius = prop_radius_list[scale_idx]

            # add position to features
            feature0, feature1 = feature_add_position(feature0, feature1, attn_splits, self.feature_channels)

            # Transformer
            feature0, feature1 = self.transformer(feature0, feature1, attn_num_splits=attn_splits)

            # correlation and softmax
            if corr_radius == -1:  # global matching
                flow_pred = global_correlation_softmax(feature0, feature1, pred_bidir_flow)[0]
            else:  # local matching
                flow_pred = local_correlation_softmax(feature0, feature1, corr_radius)[0]

            # flow or residual flow
            flow = flow + flow_pred if flow is not None else flow_pred

            # upsample to the original resolution for supervison
            if self.training:  # only need to upsample intermediate flow predictions at training time
                flow_bilinear = self.upsample_flow(flow, None, bilinear=True, upsample_factor=upsample_factor)
                flow_preds.append(flow_bilinear)

            # flow propagation with self-attn
            if pred_bidir_flow and scale_idx == 0:
                feature0 = torch.cat((feature0, feature1), dim=0)  # [2*B, C, H, W] for propagation
            flow = self.feature_flow_attn(feature0, flow.detach(),
                                          local_window_attn=prop_radius > 0,
                                          local_window_radius=prop_radius)

            # bilinear upsampling at training time except the last one
            if self.training and scale_idx < self.num_scales - 1:
                flow_up = self.upsample_flow(flow, feature0, bilinear=True, upsample_factor=upsample_factor)
                flow_preds.append(flow_up)

            if scale_idx == self.num_scales - 1:
                flow_up = self.upsample_flow(flow, feature0)
                flow_preds.append(flow_up)

        results_dict.update({'flow_preds': flow_preds})

        return results_dict

• conclusion 1: 多帧一起求光流时, 得到的光流[每帧index]对应的光流表示 当前帧 和 下一帧的光流.
• conclusion 2: 得到的光流[最后一帧index]对应的光流表示 最后一帧 到 第一帧的光流.

您提到的关于光流计算结果的形状与预期不符的问题，以及对[2*B, 2, H, W]和[2*(B-1), 2, H, W]两个形状的疑惑，主要涉及光流计算中帧对的处理方式。在此，我将解释为何计算的双向光流shape为[2*B, 2, H, W]而非[2*(B-1), 2, H, W]，以及额外的光流表示什么。

1. 双向光流计算

在视频处理中，双向光流通常指的是从帧i到帧i+1的前向光流和从帧i+1到帧i的后向光流。对于给定的一组连续帧，若我们有B个帧，则可以计算B-1对相邻帧之间的双向光流。然而，在实际应用中，为了保持计算的一致性和便于处理，通常会采用一种称为“周期边界条件”（Periodic Boundary Condition）的方法来处理首尾帧。

2. 周期边界条件

周期边界条件假设视频序列是循环的，即最后一帧之后紧跟着第一帧，第一帧之前则是最后一帧。因此，除了计算B-1对相邻帧之间的光流外，还会额外计算帧B（最后一帧）到帧1（第一帧）以及帧1到帧B的光流。这样做的好处在于：

完整性：确保整个视频序列中的所有帧都参与了光流计算，避免了边界帧信息的丢失。
一致性：在进行某些基于光流的操作（如视频插帧、稳定化等）时，周期边界条件可以避免因边界处理不当导致的不连续或伪影。

3. 形状解释

根据上述分析，对于B个帧，我们实际计算了B-1对相邻帧的双向光流，再加上首尾帧之间的一对光流，总计B对双向光流。每一对双向光流包含前向光流和后向光流，每个光流的形状为[2, H, W]（分别对应水平和垂直方向的位移）。

可视化

gmflow中采用矩阵乘法求两个图像的全局匹配，而以前的光流算法多采用一定半径的局部求相关性，然而GMFlow进行两个超大矩阵的乘法会严重增加计算量！这里对全局匹配的结果进行可视化，进行直观感受。

GMFlow中全局匹配的Code：就是对两个特征用matmul矩阵乘法，求相关性。后面可视化就是对prob的可视化

# global correlation
b, c, h, w= feature0.shape
feature0 = feature0.view(b, c, -1).permute(0, 2, 1)  # [B, H*W, C]
feature1 = feature1.view(b, c, -1)  # [B, C, H*W]

correlation = torch.matmul(feature0, feature1).view(b, h, w, h, w) / (c ** 0.5)  # [B, H, W, H, W]
correlation = correlation.view(b, h * w, h * w)  # [B, H*W, H*W]
prob = F.softmax(correlation, dim=-1)  # [B, H*W, H*W]

CAM的工作原理：

• 由于整个模型在训练时，会发生变化的参数只有分类器的权重矩阵 $W$，因此对于同一张图片，卷积层输出的特征图集始终不变，但分类概率会随着 $W$ 的变化而不断改变，这也就是模型学习的过程，即：到底哪些特征图对于提高分类准确率最有帮助？

• 可以看到，模型做出分类决策的依据来源于 $W$ 矩阵。那么如何进行可视化呢？ $W$ 矩阵本身并不直观。这里我们可以注意到，由于 $W$ 矩阵对图像的理解是基于对特征向量的加权，而特征向量背后是一个个特征图，因此可以跳过特征向量，直接将这些特征图用 $W$ 加权，再重叠合成为一张特征图，就可以很直观的看到到底模型是通过看哪片区域来做出判断的。

• 举例来说，对于一张猫咪的图片，刚开始训练时，$W$ 的值还是初始化的值，此时模型听信 $W$ 矩阵，选择了第10/20/30号特征图作为判断依据，判断图像 84% 是狗。由于这么预测的Loss很大，因此在后续反向传播的时候不断更新 $W$ 矩阵，开始逐渐以第99/100/101号等特征图作为判断依据，直到能够判断图像 100% 是猫的时候，模型基本上就学会了如何判断猫狗。

CAM可视化整体流程：

• 归一化：由于ColorMap的工作原理是将任意矩阵的取值范围映射到0~255范围内，因此为了之后好挑选颜色，需要归一化一下。
• 手动放大至[0,255]
• 元素格式化
• 设定门限：由于只希望将模型认为比较重要的区域标记出来，因此应该选择一个门限值，将该门限值之下的像素置 0（将该像素透明化）
• 将源图像与热图合成为一张图
• 在图像本体添加文字标注（方便之后导出为视频）

# 根据卷积层输出特征图集和模型某一参数状态计算预测概率（为了简单省略了bias计算）
def predict_on_weights(out_base, weights):
    gap = np.average(out_base, axis=(0, 1))
    logit = np.dot(gap, np.squeeze(weights))
    return 1 / (1 +  np.e ** (-logit))


def getCAM(image, feature_maps, weights, display=False):
    predict = predict_on_weights(feature_maps, weights)
    
    # Weighted Feature Map
    cam = (predict - 0.5) * np.matmul(feature_maps, weights)
    # Normalize
    cam = (cam - cam.min()) / (cam.max() - cam.min())
    # Resize as image size
    cam_resize = cv2.resize(cam, (224, 224))
    # Format as CV_8UC1 (as applyColorMap required)
    cam_resize = 255 * cam_resize
    cam_resize = cam_resize.astype(np.uint8)
    # Get Heatmap
    heatmap = cv2.applyColorMap(cam_resize, cv2.COLORMAP_JET)
    # Zero out
    heatmap[np.where(cam_resize <= 100)] = 0
    
    out = cv2.addWeighted(src1=image, alpha=0.8, src2=heatmap, beta=0.4, gamma=0)
    out = cv2.resize(out, dsize=(400, 400))
    
    if predict < 0.5:
        text = 'cat %.2f%%' % (100 - predict * 100)
    else:
        text = 'dog %.2f%%' % (predict * 100)
        
    cv2.putText(out, text, (210, 40), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.9, 
                  color=(123,222,238), thickness=2, lineType=cv2.LINE_AA)
    if display:
        plt.figure(figsize=(7, 7))
        plt.imshow(out[:, :, ::-1])
        plt.show()
    return out

out_base = base_model.predict(np.expand_dims(target, axis=0))
out_base = out_base[0]
print(out_base.shape)
getCAM(image=target, feature_maps=out_base, weights=weights_history[1249], display=True);

GMFlow的CAM可视化整体流程：

• 在feature0的尺度上，H高度，W宽度方向分别均匀采样N=20个点，整幅图就有$N^2$个点
• 画出每个点的Class Activation Map(CAM)图。
• 最后对所有点的CAM图求和，再和原图相加得到最后的可视化图像。

可以看到CAM可视化结果：

• 相关性上还是呈现一个局部相关性，所有用matmul理论上应该是非必须的。可以以当前点为基点，在一定范围内做局部匹配
• 核心的光流部分点的匹配，相较于静止点，其AM还是有不一样的。