MAST: A Memory-Augmented Self-Supervised Tracker论文解读和代码剖析

官方代码
作者开源的官方代码有一处错误，在代码剖析部分将指出。有人已经在github上提出了issue，作者一直没回应。我也是在阅读代码的时候发现了这个错误。

背景

VOS任务很少有使用自监督的，即在训练中不借助mask，只用frame image来训练。
作者巧妙的在STM的基础上，将value换成frame自身，使用过去帧重构当前帧作为代理任务（proxy），实现自监督的vos。效果还不错，在davis val上是64的J&F。

核心思想

仍然是采用STM的memory bank的思想。memory的特征和query的特征会使用transform，得到attention map。但不同的是，stm使用的是经过backbone得到的value，而MAST是直接使用raw frame或者mask。如果是训练阶段；使用raw frame，如果是test阶段，直接使用得到的mask。
在训练阶段，使用当前帧的特征作为query，和memory中的key，value是对应时刻的raw frame，直接使用qkv三元组重构出一个新的帧。这个输出又以当前帧为GT，用huber loss优化。整个过程没有使用到mask GT。在测试阶段，直接使用mask代替raw frame，则每次预测得到的都是重构出来的mask，作为当前帧的输出。

细节

颜色空间

作者认为，RGB颜色空间不适合作为输入，因为是重构作为代理任务。用huber loss是直接优化像素距离的。
比如说重构出来的输出的第i个像素，和raw frame的第i 个像素的matching 距离很小，但实际上他们可能是落在不同目标上。则说明，根据颜色匹配来优化网络，不适合推动模型学习语义特性。
作者也否决了随意丢弃一个channel的做法，因为RGB是关联的，可以通过其他两个通道推理得知另一个通道的像素。
作者使用LAB空间，在随机丢弃一个channel。lab空间解耦性较好。

作者统计了davis数据集的RGB数值和LAB数值的分布图。可以看出RGB是线性相关的。
输入的颜色值不是互相关联的，则网络将被push学习更好的表征，而并非仅仅依赖局部颜色信息。

loss

作为GT，raw frame使用RGB颜色空间。使用smooth l1 loss（huber loss）

 outputs = F.interpolate(outputs, (h, w), mode='bilinear')
 loss = F.smooth_l1_loss(outputs*20, tar_y*20, reduction='mean')

获取ROI区域

作者分析了STM的劣势，就是memory bank式的matching，需要的内存和计算量都很大，O(T*H*W*H*W)。
如果先获得了目标的大致位置，每一个pixel需要匹配的数目就会少很多（原始的是T*H*W)。
作者提出了一个两阶段的ROI localization。假设对query的第i个位置$q_i$进行匹配。首先使用一个网格（应用空洞技巧），围绕在key的第i个位置上，得到网络上的特征，和$q_i$做匹配（dot运算），得到的相似性系数直接加权相对坐标（和直推式vos的做法类似），这里是应用soft argmax，得到离第i个位置最相似的offset。
第二步就是围绕新的位置（i+offset），resample出一个小区域，作为需要匹配的对象。

其他细节

网络使用resnet18，修改stride，最低分辨率为1/4。训练也是先pretrain，在main train，接着dynamic train。

代码剖析

主要看看ROI那步。其他的步骤都很好读
作者先是在init里面设置了两种sampler。第一个是带dilate的，第二种是没有dilation的。前者用于long term的sampler，后者用于short term。

self.correlation_sampler_dilated = [
            SpatialCorrelationSampler(
            kernel_size=1,
            patch_size=self.memory_patch_P,
            stride=1,
            padding=0,
            dilation=1,
            dilation_patch=dirate) for dirate in range(2,6)]

self.correlation_sampler = SpatialCorrelationSampler(
      kernel_size=1,
      patch_size=self.P,
      stride=1,
      padding=0,
      dilation=1)

在forward里面，大致有下面几个步骤：

• 先对long term key进行第一步粗糙采样，得到ROI的位置，然后在截取主要特征作为matching对象得到系数。
• 在对short term key同样操作
• 用得到的offset，对raw frames，也截取对应的value。
• 所有的attention map以及value都齐了，开始使用qkv公式得到输出。

 for searching_index in range(nsearch):  # long term: need dilation
            ##### GET OFFSET HERE.  (b,h,w,2)
            samplerindex = dirates[searching_index]-2
            coarse_search_correlation = self.correlation_sampler_dilated[samplerindex](feats_t, feats_r[searching_index])  # b, p, p, h, w
            coarse_search_correlation = coarse_search_correlation.reshape(b, self.memory_patch_N, h*w)
            coarse_search_correlation = F.softmax(coarse_search_correlation, dim=1)
            coarse_search_correlation = coarse_search_correlation.reshape(b,self.memory_patch_P,self.memory_patch_P,h,w,1)
            _y, _x = torch.meshgrid(torch.arange(-self.memory_patch_R,self.memory_patch_R+1),torch.arange(-self.memory_patch_R,self.memory_patch_R+1))
            grid = torch.stack([_x, _y], dim=-1).unsqueeze(-2).unsqueeze(-2)\
                .reshape(1,self.memory_patch_P,self.memory_patch_P,1,1,2).contiguous().float().to(coarse_search_correlation.device)
            # 每个query像素在mem bank中的一帧该以哪个位置为中心采样
            offset0 = (coarse_search_correlation * grid ).sum(1).sum(1) * dirates[searching_index]  # 1,h,w,2

            col_0 = deform_im2col(feats_r[searching_index], offset0, kernel_size=self.P)  # b,c*N,h*w
            col_0 = col_0.reshape(b,c,N,h,w)
            ##
            corr = (feats_t.unsqueeze(2) * col_0).sum(1)   # (b, N, h, w)

            corr = corr.reshape([b, self.P * self.P, h * w])
            corrs.append(corr)

 for ind in range(nsearch, nref):  # short term
            corrs.append(self.correlation_sampler(feats_t, feats_r[ind]))
            _, _, _, h1, w1 = corrs[-1].size()
            corrs[ind] = corrs[ind].reshape([b, self.P*self.P, h1*w1])

得到T帧的匹配系数的softmax值

  corr = torch.cat(corrs, 1)  # b,nref*N,HW
        corr = F.softmax(corr, dim=1)
        corr = corr.unsqueeze(1)

得到value

im_col0 = [deform_im2col(qr[i], offset0, kernel_size=self.P)  for i in range(nsearch)]# b, 3*N, h*w
im_col1 = [F.unfold(r, kernel_size=self.P, padding=self.R) for r in qr[nsearch:]]
image_uf = im_col0 + im_col1  # memory value list.

得到预测结果

  out = (corr * image_uf).sum(2).reshape([b,qr[0].size(1),h,w])

采用使用的是spatial correlation sapmle，是计算光流的cost valume的重要操作。不知道啥是cost valume可以去知乎搜索一下。作者这里用他是计算$q_i$和在key上以i为中心的网格中被选取的特征的相似度。

所谓的截取，就是已知$q_i$应该在哪个位置截取，就使用grid sample取出来。