視頻目標分割之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