视频目标分割之Video Object Segmentation using Space-Time Memory Networks

Video Object Segmentation using Space-Time Memory Networks

原始文档： https://www.yuque.com/lart/papers/iw3gs6

优缺点

优点

• 简单
  • 模型：
    • 编解码结构：
      • 编码器提取特征，通过匹配网络与之前帧目标的特征进行匹配
      • 解码器使用类fpn结构，逐步恢复特征，生成最终预测
    • 没有后处理，没有在线学习，所以只要解决了测试过程中的存储需求，基本就很简单了
  • 思路：
    • 不仅仅是用第一针和前一帧，还会采样更加靠前的帧的信息用来生成当前帧的预测
    • 扩展memory network思路到了VOS中：网络中维持着一个memory单元，存储由memory network提取的之前特定帧和对应的mask（可能是softmax的输出）的特征。通过使用query nertowrk提取当前帧的特征，从memory中查询并匹配信息，最终生成预测
      • 来自于NLP中Q&A任务的一种经典结构。在这种方法中，可记忆的信息被分别嵌入到key（即，input）和value（即，output）特征向量中。key通常被用来在memory中寻址更加相关的信息（把memory看做是一种存储结构），并且它对应的value会被返回
      • 帧和对应的mask（实际除了整个视频序列、或者视频片段的第一帧使用的是原始的mask之外，其他的都是网络最终softmax处理后的概率张量）提取的特征会被存放到memory中。memory是动态更新的，会不断的添加新的之前帧的预测的mask（实际是softmax的输出）
      • 单独的当前帧会被送入query network中，生成特征、从memory检索匹配，定位并读取整合相关的信息，最终生成预测
      • memory单元的读取结构可以看作是一个时空注意力算法（类似non-local），可以确定每个query中的像素是输入前景目标还是背景，从而辅助多目标的类别确定
  • 策略：
    • 使用更多的数据：对静态数据处理生成合适的训练样本进行大规模预训练，再在视频数据上微调
    • 使用更多的参考帧：
      • 在训练的时候，使用固定数量（文中设定为3）的帧作为单个epoch中单次迭代的输入
      • 在测试的时候，使用固定间隔（文中设定为5）采样之前帧的特征用来作为memory，辅助生成当前帧的预测。注意：这个地方是显存占用的大户，因为使用的是类似non-local的结构进行的检索，矩阵乘法造成的占用巨大
• 性能非常好

缺点

• 训练
  • 步骤繁琐
  • 并没有提供训练代码，有许多细节可能需要翻查issue或者自己试探
• 测试
  • 内存占用非常大。但是可以使用一些策略降低对于内存的需求
  • 显存（主要是memory单元和查询的部分）占用也很大。显存占用很难调节，因为有些占用是必须的，很难通过修改模型之外的策略来缓解（除非有更大的显存）

主要结构

来自论文的结构图

来自补充材料的细节结构图

整体结构如上图所示，还是比较简单直接的。这里的Memory结构对应了我前面提到的memory network和memory单元（实际上我说的memory network更像是这里提到的memory encoder）。

来自论文的结构图

这是其中的memory读取的结构，也就是匹配加权整合的过程。结构很直观。

具体网络细节可以看作者提供的model代码。

处理多目标

对于多目标任务而言，多目标的处理方式也是值得关注的一个地方。

本文是通过针对不同目标分别进行独立的预测。实际上在实现中，目标数量这一维度被放在了batch的位置上，随意可以看做是等效的“独立处理”，所以实际上，对于网络的一次迭代而言，每次处理的是单独的一段采样视频（即，batchsize=1），不同帧被拼接到了另一个单独的T的维度上，一次迭代中会循环处理这些数据，可以看做是一种BPTT（RNN）的更新方式。但是论文中提到，使用的batchsize=4，这是怎么回事呢？实际是因为作者使用了四块卡，所以对于每块卡而言，处理的还是batchsize=1的数据。

对于不同的目标预测结束后，最终会进行一个针对不同目标的预测的整合策略，该整合策略集成并发展了该作者另一篇文章（Fast video object segmentation by reference-guided mask propagation，这里的mask融合策略被在测试阶段，当做一种后处理步骤，而在本文中，将整合策略调整成了可微分的操作，整合到了网络中，所以训练测试都可以使用）的思路：

这里的$\sigma$和$l$表示softmax和对应的logit函数，$\hat{p}_{i, m}$是网络将位置$i$判定为目标$m$的概率，$m=0$表示背景类别。而这里的$M$表示总的目标数量。

class STM(nn.Module):
    ...
    def Soft_aggregation(self, ps: torch.Tensor, K: int) -> torch.Tensor:
        """
        整合针对不同目标的预测

        Args:
            ps : num_objects, H, W 各通道表示像素属于各个目标的概率
            K: 所有视频里最大的目标数量

        Returns: 整合后得到的针对各个目标的预测概率
        """
        num_objects, H, W = ps.shape
        em = ToCuda(torch.zeros(1, K, H, W))
        em[0, 0] = torch.prod(1 - ps, dim=0)  # bg prob
        em[0, 1: num_objects + 1] = ps  # obj prob
        em = torch.clamp(em, 1e-7, 1 - 1e-7)
        logit = torch.log((em / (1 - em)))
        return logit
    
# in the training/testing loop:
logit = model(...)  # it will call `self.Soft_aggregation`
output_tensor = F.softmax(logit, dim=1)
# 数组形式的最终的预测
output_array = np.argmax(output_tensor[0].cpu().numpy(), axis=0)[0].astype(np.uint8)

实验细节

• 双阶段训练：
  • Pre-training on images.
    • One advantage of our framework is that it does not require long training videos. This is because the method learns the semantic spatio-temporal matching between distant pixels without any assumption on temporal smoothness. This means that we can train our network with only a few frames （最小是2，即一帧memory，一帧query） with object masks. This enables us to simulate training videos using image datasets. Some previous works [26, 24] trained their networks using static images and we take a similar strategy.
    • A synthetic video clip that consists of 3 frames is generated by applying random affine transforms (rotation, sheering, zooming, translation, and cropping) to a static image with different parameters. We leverage the image datasets annotated with object masks (salient object detection [ECSSD, MSRA10K], semantic segmentation – [PASCAL VOC, SBD, COCO]) to pre-train our network.
    • By doing so, we can expect our model to be robust against a wide variety of object appearance and categories.
  • Main training on videos.
    • After the pre-training, we use real video data for the main training. In this step, either **Youtube-VOS or DAVIS-2017 **is used, depending on the target evaluation benchmark.
    • To make a training sample, we sample 3 temporally ordered frames from a training video.
    • To learn the appearance change over a long time, we randomly skip frames during the sampling. The maximum number of frames to be skipped is gradually increased from 0 to 25 during the training as in curriculum learning [Weakly-supervised disentangling with recurrent transformations for 3d view synthesis].
• Training details.
  • We used randomly cropped 384×384 patches for training.
  • For all experiments, we set the minibatch size to 4 and disabled all the batch normalization layers.
  • We minimize the cross-entropy loss using Adam optimizer with a fixed learning rate of 1e-5.
  • Pre-training takes about 4 days and main training takes about 3 days using four NVIDIA GeForce 1080 Ti GPUs.
• Inference.
  • Writing all previous frames on to the memory may raise practical issues such as GPU memory overflow and slow running speed. Instead, we select frames to be put onto the memory by a simple rule. The first and the previous frame with object masks are the most important reference information.
    • The first frame always provides reliable information as it comes with the ground truth mask.
    • The previous frame is similar in appearance to the current frame, thus we can expect accurate pixel matching and mask propagation.
    • Therefore, we put these two frames into the memory by default.
  • For the intermediate frames, we simply save a new memory frame every N frames. N is a hyperparameter that controls the trade-off between speed and accuracy, and we use N = 5 unless mentioned otherwise.
  • It is noteworthy that our framework achieves the effect of online learning and online adaptation without additional training. The effect of online model updating is easily accomplished by putting the previous frames into the memory without updating model parameters. Thus, our method runs considerably faster than most of the previous methods while achieving state-of-theart accuracy.
• We evaluate our model on Youtube-VOS and DAVIS benchmarks. We prepared two models trained on each benchmarks’ training set.
  • For the evaluation on Youtube-VOS, we used 3471 training videos following the official split.
  • For DAVIS, we used 60 videos from the DAVIS-2017 train set. Both DAVIS 2016 and 2017 are evaluated using a single model trained on DAVIS-2017 for a fair comparison with the previous works [24, 40].
  • In addition, we provide the results for the DAVIS with our model trained with additional training data from Youtube-VOS. Note that we use the network output directly without postprocessing to evaluate our method.

论文提供了丰富的对比实验，但是除了消融试验中提到了没有使用静态数据预训练的模型性能外，其他地方都是使用预训练后的微调结果。

这可以看出来，这里的预训练是多么的有用！VOS中的数据量还是太少，这种预训练策略感觉可以看做是STM成功的主要原因。另外，这里也展示分别使用YoutubeVOS和DAVIS数据训练后，再在对方的数据上评估的结果，可以看到，还是反映了数据量问题。DAVIS本身数据量还是太小。

这里展示了测试时的采样策略，可以看到，每5帧抽取一次是很有效的。但是第一帧和前一帧同时使用带来的增益还是极大的。

相关链接

• 代码：
  • 作者（没有提供train代码）：https://github.com/seoungwugoh/STM
  • 第三方实现（提供了完整的trian和test代码）：https://github.com/lyxok1/STM-Training
• 论文：https://openaccess.thecvf.com/content_ICCV_2019/papers/Oh_Video_Object_Segmentation_Using_Space-Time_Memory_Networks_ICCV_2019_paper.pdf