DeAOT视频追踪论文阅读笔记

2023-08-08
Author AsanoSaki
~14.65K words

1. 相关知识
2. 半监督VOS与AOT模型
3. DeAOT
1. 3.1 分层双分支传播
2. 3.2 门控传播模块GPM

本文记录 DeAOT 视频追踪论文的阅读笔记。
涉及的相关知识点为：AOT（Associating Objects with Transformers for Video Object Segmentation）、DeAOT（Decoupling Features in Hierarchical Propagation for Video Object Segmentation）、FPN（Feature Pyramid Networks for Object Detection）、Depth-wise Convolution、DropPath、GroupNorm。

1. 相关知识

1.1 深度可分离卷积

Depth-wise（DW）卷积与 Point-wise（PW）卷积，合起来被称作 Depth-wise Separable Convolution（深度可分离卷积），该结构和常规卷积操作类似，可用来提取特征，但相比于常规卷积操作，其参数量和运算成本较低。所以在一些轻量级网络中会碰到这种结构，如 MobileNet。

Depth-wise Convolution 的一个卷积核负责一个通道，即一个通道只被一个单通道的卷积核卷积，而常规卷积每个卷积核是同时操作输入图片的每个通道，即每个卷积核的通道数与图片的通道数相同。

Depth-wise Convolution 完成后的 Feature Map 数量与输入层的通道数相同，无法扩展 Feature Map。而且这种运算对输入层的每个通道独立进行卷积运算，没有有效地利用不同通道在相同空间位置上的特征信息。因此需要Point-wise Convolution 来将这些 Feature Map 进行组合生成新的 Feature Map。

Depth-wise Convolution 代码如下：

import torch
import torch.nn as nn

# Depth-wise卷积，输出维度和输入维度相同
class DWConv2d(nn.Module):
    def __init__(self, in_channels, dropout=0.1):
        super().__init__()
        # 当groups=in_channels时，是在做Depth-wise Conv
        self.conv = nn.Conv2d(in_channels, in_channels, 3, padding=1, groups=in_channels)
        self.dropout = nn.Dropout2d(p=dropout, inplace=True)

    def forward(self, x):  # x.shape: (bsz, c, h, w)
        out = self.dropout(self.conv(x))
        return out  # out.shape: (bsz, c, h, w)

dwconv2d = DWConv2d(in_channels=3)
x = torch.randn(1, 3, 10, 10)
print(dwconv2d(x).shape)  # torch.Size([1, 3, 10, 10])

Point-wise Convolution 的运算与常规卷积运算非常相似，它的卷积核的尺寸为 1 × 1 × M，M 为输入图像的通道数。所以这里的卷积运算会将上一步的 Feature Map 在深度方向上进行加权组合，生成新的 Feature Map。有几个卷积核就有几个输出 Feature Map。

Point-wise Convolution 代码如下：

# Point-wise卷积，输出维度可以改变
class PWConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, dropout=0.1):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, 1)
        self.dropout = nn.Dropout2d(p=dropout, inplace=True)

    def forward(self, x):  # x.shape: (bsz, c, h, w)
        out = self.dropout(self.conv(x))
        return out  # out.shape: (bsz, c', h, w)

pwconv2d = PWConv2d(in_channels=3, out_channels=4)
x = torch.randn(1, 3, 10, 10)
print(pwconv2d(x).shape)  # torch.Size([1, 4, 10, 10])

Depth-wise 卷积对每个输入通道独立执行空间卷积，而 Point-wise 卷积用于组合 Depth-wise 卷积的输出。将两者组合起来即为深度可分离卷积神经网络：

class DepthwiseSeparableConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, dropout=0.1):
        super().__init__()
        self.dwconv2d = DWConv2d(in_channels, dropout)
        self.pwconv2d = PWConv2d(in_channels, out_channels, dropout)

    def forward(self, x):  # x.shape: (bsz, c, h, w)
        out = self.pwconv2d(self.dwconv2d(x))
        return out  # out.shape: (bsz, c', h, w)

depthwiseseparableconv2d = DepthwiseSeparableConv2d(3, 4)
x = torch.randn(1, 3, 10, 10)
print(depthwiseseparableconv2d(x).shape)  # torch.Size([1, 4, 10, 10])

1.2 DropPath

DropPath 是一种针对分支网络而提出的网络正则化方法，作用是将深度学习网络中的多分支结构随机删除。DropPath 作用的是网络分支，而 DropOut 作用的是 Feature Map，DropConnect 作用的是参数。

简单来说，DropPath 的输出是随机将一个 batch 中所有的神经元均设置为0；而在 DropOut 中，是在每个 batch 中随机选择神经元设置为0。

DropPath 代码如下：

import torch
import torch.nn as nn

class DropPath(nn.Module):
    def __init__(self, drop_prob=None, batch_dim=0):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob  # 丢弃率，假设是0.5
        self.batch_dim = batch_dim  # batch在第几维，假设是0

    def forward(self, x):
        return self.drop_path(x, self.drop_prob)

    def drop_path(self, x, drop_prob):  # x.shape: (hw, bsz, c)
        # 丢弃率为0或者不是在训练时直接返回x
        if drop_prob == 0. or not self.training:
            return x

        keep_prob = 1 - drop_prob  # 保持率，0.5
        shape = [1 for _ in range(x.ndim)]  # [1, 1, 1]
        shape[self.batch_dim] = x.shape[self.batch_dim]  # [bsz, 1, 1]
        random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)  # 0~1之间的均匀分布
        random_tensor.floor_()  # 下取整，随机出来的大于等于0.5的数都为1，确定保留哪些batch
        output = x.div(keep_prob) * random_tensor  # 除以keep_prob是为了让训练和测试时的期望保持一致

        return output

droppath = DropPath(drop_prob=0.5, batch_dim=0)
a = torch.randn(3, 2, 3)
print(droppath(a))
# tensor([[[-0.6209, -5.4889, -1.9857],
#          [ 0.1626,  6.0644,  0.8875]],
#
#         [[ 0.0000, -0.0000,  0.0000],
#          [ 0.0000, -0.0000,  0.0000]],
#
#         [[-0.1041,  0.4921,  0.3389],
#          [-2.0490, -0.0399, -0.1521]]])

1.3 Group Normalization

BN 全名是 Batch Normalization，见名知意，其是一种归一化方式，而且是以 batch 的维度做归一化，那么问题就来了，此归一化方式如果使用过小的 batch size 会导致其性能下降，一般来说每个 GPU 上的 batch size 设为32最合适，但是对于一些其他深度学习任务 batch size 往往只有1或2，比如目标检测，图像分割，视频分类上，输入的图像数据很大，较大的 batch size 显存吃不消。

另外，Batch Normalization 是在 batch 这个维度上做 Normalization，但是这个维度并不是固定不变的，比如训练和测试时一般不一样，一般都是训练的时候在训练集上通过滑动平均预先计算好平均（mean），和方差（variance）参数，在测试的时候，不再计算这些值，而是直接调用这些预计算好的参数来用。但是，当训练数据和测试数据分布有差别时，训练机上预计算好的数据并不能代表测试数据，这就导致在训练、验证、测试这三个阶段存在 inconsistency（不一致性）。

Group Normalization（GN）首先将 channel 分为许多组（group），对每一组做归一化，即先将 feature 的维度由 [N, C, H, W] reshape 为 [N, G, C/G, H, W]，归一化的维度为 [C/G, H, W]。事实上，GN 的极端情况就是 LN（Layer Normalization）和 IN（Instance Normalization），分别对应 G = 1 和 G = C，作者在论文中给出 G 设为32较好。

Group Normalization 代码如下：

import torch
from torch import nn

# 自己实现GN
class GroupNorm(nn.Module):
    def __init__(self, num_channels, num_groups=32, eps=1e-5):
        super(GroupNorm, self).__init__()
        self.gamma = nn.Parameter(torch.ones(1, num_groups, 1))
        self.beta = nn.Parameter(torch.zeros(1, num_groups, 1))
        self.num_groups = num_groups
        self.eps = eps

    def forward(self, x):  # x.shape: (N, C, H, W)
        N, C, H, W = x.size()
        G = self.num_groups
        assert C % G == 0

        x = x.view(N, G, -1)
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)

        x = self.gamma * (x - mean) / (std + self.eps) + self.beta
        x = x.view(N, C, H, W)

        return x

x = torch.randn(1, 2, 2, 3)
gn = GroupNorm(2, 2)
x_gn = gn(x)

print(x_gn)
# tensor([[[[-0.1165, -0.5803,  0.7635],
#           [-0.2374, -1.3281,  1.4988]],
#
#          [[ 0.0982, -0.3936, -1.3941],
#           [-0.5678,  1.2321,  1.0253]]]], grad_fn=<ViewBackward0>)
print(x_gn.mean((2, 3)))  # tensor([[-4.9671e-09,  9.9341e-09]], grad_fn=<MeanBackward1>)
print(x_gn.var((2, 3)))  # tensor([[1.0000, 1.0000]], grad_fn=<VarBackward0>)

# 使用torch.nn中的GN
torch_gn = nn.GroupNorm(num_groups=2, num_channels=2)
x_torch_gn = torch_gn(x)

print(x_torch_gn)
# tensor([[[[-0.1277, -0.6357,  0.8363],
#           [-0.2601, -1.4548,  1.6419]],
#
#          [[ 0.1076, -0.4312, -1.5272],
#           [-0.6220,  1.3497,  1.1232]]]], grad_fn=<NativeGroupNormBackward0>)
print(x_torch_gn.mean((2, 3)))  # tensor([[-2.9802e-08,  1.9868e-08]], grad_fn=<MeanBackward1>)
print(x_torch_gn.var((2, 3)))  # tensor([[1.2000, 1.2000]], grad_fn=<VarBackward0>)

1.4 FPN特征金字塔网络

目标的多尺度一直是目标检测算法极为棘手的问题。像 Fast R-CNN，YOLO 这些只是利用深层网络进行检测的算法，是很难把小目标物体检测好的。因为小目标物体本身的像素就比较少，随着下采样的累积，它的特征更容易被丢失。

特征金字塔网络（Feature Pyramid Network，FPN）是一个在特征尺度的金字塔操作，它是通过将自底向上（Bottom-up）和自顶向下（Top-down）的特征图进行融合来实现特征金字塔操作的。FPN 提供的是一个特征融合的机制，并没有引入太多的参数，实现了以增加极小计算代价的情况下提升对多尺度目标的检测能力。

自底向上即是卷积网络的前向过程，我们可以选择不同的骨干网络，例如 ResNet-50 或者 ResNet-101。前向网络的返回值依次是 C2、C3、C4、C5，是每次池化之后得到的 Feature Map。

通过自底向上路径，FPN 得到了四组 Feature Map。浅层的 Feature Map，例如 C2 含有更多的底层信息（纹理，颜色等），而深层的 Feature Map 如 C5 含有更多的语义信息。为了将这四组倾向不同特征的 Feature Map 组合起来，FPN 使用了自顶向下及横向连接的策略，最终得到 P2、P3、P4、P5 四个输出。

最后，FPN 在 P2、P3、P4、P5 之后均接了一个 3*3 Conv 操作，该卷积操作是为了减轻上采样的混叠效应（aliasing effect）。

FPN 和 U-Net 最大的不同是它的多个层级的都会有各自的输出层，而每个输出层都有不同尺度的感受野。一个比较粗暴的方式是每一层都预测所有的样本，而另一个更好的选择是根据一些可能存在的先验知识选择一个最好的层。

FPN 代码如下：

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class Bottleneck(nn.Module):
    expansion = 4  # 残差块第3个卷积层的通道膨胀倍率

    def __init__(self, in_channels, channels, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, channels, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(channels)
        self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(channels)
        self.conv3 = nn.Conv2d(channels, self.expansion * channels, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(self.expansion * channels)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x  # 将原始输入暂存为shortcut的输出
        if self.downsample is not None:  # 如果需要下采样，那么shortcut后：H/2，W/2。C：out_channel -> 4 * out_channel
            residual = self.downsample(x)

        out = self.relu(self.bn1(self.conv1(x)))
        out = self.relu(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))

        out += residual  # 残差连接
        out = self.relu(out)

        return out

class FPN(nn.Module):
    def __init__(self, block, layers):
        super(FPN, self).__init__()
        self.inchannels = 64

        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(64)

        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

        # Bottom-up layers
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)

        # Top layer
        self.toplayer = nn.Conv2d(2048, 256, kernel_size=1, stride=1, padding=0)  # Reduce channels

        # Lateral layers
        self.latlayer1 = nn.Conv2d(1024, 256, kernel_size=1, stride=1, padding=0)
        self.latlayer2 = nn.Conv2d(512, 256, kernel_size=1, stride=1, padding=0)
        self.latlayer3 = nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)

        # Smooth layers
        self.smooth1 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
        self.smooth2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
        self.smooth3 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self, block, channel, block_num, stride=1):
        downsample = None
        if stride != 1 or self.inchannels != channel * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inchannels, block.expansion * channel, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(block.expansion * channel)
            )
        layers = []
        layers.append(block(self.inchannels, channel, stride, downsample))
        self.inchannels = channel * block.expansion
        for i in range(1, block_num):
            layers.append(block(self.inchannels, channel))

        return nn.Sequential(*layers)

    def _upsample_add(self, x, y):  # 将x上采样成y的size后与y相加
        _, _, H, W = y.size()
        return F.interpolate(x, size=(H, W), mode='bilinear') + y

    def forward(self, x):  # (bsz, 3, h, w)
        # Bottom-up
        x = self.conv1(x)  # (bsz, 64, h/2, w/2)
        x = self.bn1(x)
        x = self.relu(x)
        c1 = self.maxpool(x)  # (bsz, 64, h/4, w/4)

        c2 = self.layer1(c1)  # (bsz, 256, h/4, w/4)
        c3 = self.layer2(c2)  # (bsz, 512, h/8, w/8)
        c4 = self.layer3(c3)  # (bsz, 1024, h/16, w/16)
        c5 = self.layer4(c4)  # (bsz, 2048, h/32, w/32)

        # Top-down
        p5 = self.toplayer(c5)  # (bsz, 256, h/32, w/32)
        p4 = self._upsample_add(p5, self.latlayer1(c4))  # (bsz, 256, h/16, w/16)
        p3 = self._upsample_add(p4, self.latlayer2(c3))  # (bsz, 256, h/8, w/8)
        p2 = self._upsample_add(p3, self.latlayer3(c2))  # (bsz, 256, h/4, w/4)

        # Smooth
        p4 = self.smooth1(p4)  # (bsz, 256, h/16, w/16)
        p3 = self.smooth2(p3)  # (bsz, 256, h/8, w/8)
        p2 = self.smooth3(p2)  # (bsz, 256, h/4, w/4)
        return p2, p3, p4, p5

def FPN101():
    return FPN(Bottleneck, [2, 2, 2, 2])

fpn_101 = FPN101()

input = torch.randn(1, 3, 256, 256)
output_p2, output_p3, output_p4, output_p5 = fpn_101(input)
print(output_p2.shape)  # torch.Size([1, 256, 64, 64])
print(output_p3.shape)  # torch.Size([1, 256, 32, 32])
print(output_p4.shape)  # torch.Size([1, 256, 16, 16])
print(output_p5.shape)  # torch.Size([1, 256, 8, 8])

2. 半监督VOS与AOT模型

2.1 VOS与AOT简介

视频对象分割（VOS）旨在识别和分割给定视频中的一个或多个感兴趣的对象，半监督 VOS 需要算法在给定一帧或多帧的对象注释掩码的情况下跟踪和分割整个视频序列中的对象。

此前最先进的方法学习用单个正目标解码特征，因此必须在多目标场景下单独匹配和分割每个目标，消耗多倍的计算资源。我们提出 Associating Objects with Transformers（AOT）方法来统一匹配和解码多个对象。AOT 采用 Identification 机制将多个目标关联到同一高维嵌入空间中。因此可以像处理单个对象一样高效地同时处理多个对象的匹配和分割解码。

AOT 方法将分层传播引入到 VOS 中。分层传播可以逐渐将 ID 信息从过去的帧传播到当前帧，并将当前帧的特征从 object-agnostic（对象不可知）转移到 object-specific（对象特定）。

2.2 ID 机制

ID 机制为每个目标分配唯一的 ID 信息，并将任意数量（要求小于预定义的大量）目标的 mask 嵌入到同一高维空间中。因此，网络可以学习所有目标之间的关联或相关性。此外，可以利用分配的 ID 信息直接解码多对象分割。

我们初始化一个身份库（ID Bank），其中存储 M 个具有 C 维的识别向量。为了嵌入多个不同的目标掩码，每个目标将被随机分配一个不同的识别向量。

2.3 Long Short-Term Transformer（LSTT）

本文设计长短期 Transformer（LSTT）用于构建分层对象匹配和传播。每个 LSTT 块都利用长期注意力来匹配第一帧的嵌入，并利用短期注意力来匹配多个附近帧的嵌入。与仅利用一个注意力层的方法相比，我们发现分层注意力结构在关联多个对象方面更有效。

LSTT 首先采用自注意力层，负责学习当前帧内目标之间的关联或相关性。此外，LSTT 还引入了长期注意力和短期注意力，前者用于聚合来自长期记忆帧的目标信息，后者能够从邻近的短期帧学习时间平滑性。所有注意力模块都是以多头注意力的形式实现的，即多个注意力模块后跟串联和线性投影。

长期注意力负责将目标的信息从过去的记忆帧（包含参考帧和存储的预测帧）聚合到当前帧。由于当前帧和过去帧之间的时间间隔是可变的并且可以是长期的，因此时间平滑性难以保证。因此，长期注意力采用非局部注意力。

短期注意力用于聚合每个当前帧位置的时空邻域中的信息。直观上，多个连续视频帧之间的图像变化始终是平滑且连续的。因此，连续帧中的目标匹配和传播可以限制在小的时空邻域内，从而比非局部过程具有更好的效率。

3. DeAOT

本文重点是为半监督视频对象分割（VOS）开发一种更有效的分层传播方法。在 AOT 方法中 object-specific 信息的增加将不可避免地导致深层传播层中 object-agnostic 的视觉信息的丢失。为了解决这样的问题并进一步促进视觉嵌入的学习，本文提出了一种分层传播中的解耦特征（DeAOT）方法。DeAOT 通过在两个独立的分支中处理 object-agnostic 和 object-specific 的嵌入来解耦它们的分层传播。其次，为了补偿双分支传播的额外计算，设计了一种用于构造分层传播的有效模块 GPM（门控传播模块），它是通过单头注意力精心设计的。

3.1 分层双分支传播

DeAOT 在两个并行分支中传播对象的视觉特征（visual features）和掩码（mask features）特征。具体来说，视觉分支负责匹配对象、收集过去的视觉信息并提炼对象特征。为了重新识别对象，ID 分支重用视觉分支计算的匹配图（注意力图），将 ID 嵌入（由 AOT 中的 ID 机制编码）从过去的帧传播到当前帧。两个分支共享相同的具有 L 个传播层的层次结构。

3.2 门控传播模块GPM

门控传播函数首先通过使用条件门来增强基于注意力的传播，此外，我们利用 Depth-wise 卷积以轻量级方式增强局部空间上下文的建模。

门控传播模块由三种门控传播组成：自传播、长期传播、短期传播。与 LSTT 相比，GPM 去掉了前馈模块，进一步节省了计算量和参数。所有传播过程都采用门控传播函数。

AsanoSaki