像教女朋友一樣的Deformable DETR論文精度+代碼詳解（1）

最近學(xué)習(xí)CV中的Transformer有感而發(fā)，網(wǎng)上關(guān)于Deformable DETR通俗的帖子不是很多，因此想分享一下最近學(xué)習(xí)的內(nèi)容。第一次發(fā)帖經(jīng)驗(yàn)不足，文章內(nèi)可能有許多錯誤或不恰之處歡迎批評指正。

DETR消除了目標(biāo)檢任務(wù)中的手工設(shè)計(jì)痕跡，但是存在收斂慢以及Transformer的自注意力造成的特征圖分辨率不能太高的問題，這就導(dǎo)致了小目標(biāo)檢測性能很差。我們的Deformable DETR只在參考點(diǎn)附近采樣少量的key來計(jì)算注意力，因此我們的方法收斂快并且可以用到多尺度特征。

傳統(tǒng)目標(biāo)檢測任務(wù)有很多手工設(shè)計(jì)痕跡，所以不是端到端的網(wǎng)絡(luò)。DETR運(yùn)用到了Transformer強(qiáng)大的功能以及全局關(guān)系建模能力來取代目標(biāo)檢測中人工設(shè)計(jì)痕跡來達(dá)到端到端的目的。

DETR的兩大缺點(diǎn)：

(1)收斂速度慢：因?yàn)槿窒袼刂g計(jì)算注意力要收斂到幾個稀疏的像素點(diǎn)需要消耗很長的時間。

(2)小目標(biāo)檢測差：目標(biāo)檢測基本都是在大分辨率的特征圖上進(jìn)行小目標(biāo)的檢測，但是Transformer中的Self Attention的計(jì)算復(fù)雜度是平方級別的，所以只能利用到最后一層特征圖。

可變形卷積DCN是一種注意稀疏空間位置很好的機(jī)制，但是其缺乏元素之間關(guān)系的建模能力。

綜上所述，Deformable Attention模塊結(jié)合了DCN稀疏采樣能力和Transformer的全局關(guān)系建模能力。這個模塊可以聚合多尺度特征，不需要FPN了，我們用這個模塊替換了Transformer Encoder中的Multi-Head Self- Attention模塊和Transformer Decoder中的Cross Attention模塊。

Deformable DETR的提出可以幫助探索更多端到端目標(biāo)檢測的探索。提出了bbox迭代微調(diào)策略和兩階段方法，其中iterative bounding box refinement類似Cascade R-CNN方法，two stage類似RPN。

Transformer中包含了多頭自注意力和交叉注意力機(jī)制，其中多頭自注意力機(jī)制對key的數(shù)量很敏感，平方級別的復(fù)雜度導(dǎo)致不能有太多的key，解決方法主要可以分為三類。

（1）第一類解決方法為在key上使用預(yù)定義稀疏注意力模式，例如將注意力限制在一個固定的局部窗口上，這將導(dǎo)致失去了全局信息。

（2）第二類是通過數(shù)據(jù)學(xué)習(xí)到相關(guān)的稀疏注意力。

（3）第三類是尋找自注意力中低等級的屬性，類似限制關(guān)鍵元素的尺寸大小。

圖像領(lǐng)域的注意力方法大多數(shù)都局限于第一種設(shè)計(jì)方法，但是因?yàn)閮?nèi)存模式原因速度要比傳統(tǒng)卷積慢3倍（相同的FLOPs下）。DCN可以看作是一種自注意力機(jī)制，它比自注意力機(jī)制更加高效有效，但是其缺少元素關(guān)系建模的機(jī)制。我們的可變形注意力模塊來源于DCN，并且屬于第二類注意力方法。它只關(guān)注從q特征預(yù)測得到的一小部分固定數(shù)量的采樣點(diǎn)。

目標(biāo)檢測任務(wù)一個難點(diǎn)就是高效的表征不同尺度下的物體?，F(xiàn)在有的方法比如FPN，PA-FPN，NAS-FPN，Auto-FPN，BiFPN等。我們的多尺度可變形注意力模塊可以自然的融合基于注意力機(jī)制的多尺度特征圖，不需要FPN了。

該模塊計(jì)算復(fù)雜度為: , 其中  代表特征圖維度,  和  均為圖片中的像素(pixel), 因此有  。所以計(jì)算復(fù)雜度可以簡化為  , 可以得出其與圖片像素的數(shù)量成平方級別的計(jì)算復(fù)雜度。

DETR在目標(biāo)檢測領(lǐng)域中引入了Transformer結(jié)構(gòu)并且取得了不錯的效果。這套范式摒棄了傳統(tǒng)目標(biāo)檢測中的anchor和post processing 機(jī)制，而是先預(yù)先設(shè)定100個object queries然后進(jìn)行二分圖匹配計(jì)算loss。其具體流程圖(pipeline)如下

圖1. DETR Pipeline

1、輸入圖片3×800×1066的一張圖片，經(jīng)過卷積神經(jīng)網(wǎng)絡(luò)提取特征，長寬32倍下采樣后得到2048×25×34，然后通過一個1×1 Conv進(jìn)行降維最終得到輸出shape為256×25×34.

2、positional encoding為絕對位置編碼，為了和特征完全匹配形狀也為256×25×34，然后和特征進(jìn)行元素級別的相加后輸入到Transformer Encoder中。

3、輸入到Encoder的尺寸為(25×34)×256=850×256，代表有256個token每個token的維度為850，Encoder不改變輸入的Shape。

4、Encoder的輸出和object queries輸入到Decoder中形成cross attention，object queries的維度設(shè)置為anchor數(shù)量×token數(shù)量。

5、Decoder輸出到FFN進(jìn)行分類和框定位，其中FFN是共享參數(shù)的。

tips: 雖然DETR沒有anchor，但是object queries其實(shí)就是起到了anchor的作用。

DETR缺點(diǎn)在于：

（1）計(jì)算復(fù)雜度的限制導(dǎo)致不能利用大分辨率特征圖，導(dǎo)致小目標(biāo)性能差。

（2）注意力權(quán)重矩陣往往都很稀疏，DETR計(jì)算全部像素的注意力導(dǎo)致收斂速率慢。

圖2. Deformable Attention Module

Deformable Attention Module主要思想是結(jié)合了DCN和自注意力, 目的就是為了通過在輸入特征圖上的參考點(diǎn)(reference point)附近只采樣少數(shù)點(diǎn)(deformable detr設(shè)置為3個點(diǎn))來作為注意力的  。因此要解決的問題就是：（1）確定reference point。（2）確定每個reference point的偏移量 (offset)。(3) 確定注意力權(quán)重矩陣  。在Encoder和Decoder中實(shí)現(xiàn)方法不太一樣, 加下來詳細(xì)敘述。

在Encoder部分, 輸入的Query Feature  為加入了位置編碼的特征圖(src+pos), value  的計(jì)算方法只使用了src而沒有位置編碼(value_proj函數(shù))。

（1）reference point確定方法為用了torch.meshgrid方法，調(diào)用的函數(shù)如下(get_reference_points)，有一個細(xì)節(jié)就是參考點(diǎn)歸一化到0和1之間，因此取值的時候要用到雙線性插值的方法。而在Decoder中，參考點(diǎn)的獲取方法為object queries通過一個nn.Linear得到每個對應(yīng)的reference point。

def get_reference_points(spatial_shapes, valid_ratios, device):
    reference_points_list = []
    for lvl, (H_, W_) in enumerate(spatial_shapes):
        # 從0.5到H-0.5采樣H個點(diǎn)，W同理 這個操作的目的也就是為了特征圖的對齊
        ref_y, ref_x = torch.meshgrid(torch.linspace(0.5, H_ - 0.5, H_, dtype=torch.float32, device=device),
                                        torch.linspace(0.5, W_ - 0.5, W_, dtype=torch.float32, device=device))
        ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, lvl, 1] * H_)
        ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, lvl, 0] * W_)
        ref = torch.stack((ref_x, ref_y), -1)
        reference_points_list.append(ref)
    reference_points = torch.cat(reference_points_list, 1)
    reference_points = reference_points[:, :, None] * valid_ratios[:, None]
    return reference_points

（2）計(jì)算offset的方法為對  過一個nn.Linear，得到多組偏移量，每組偏移量的維度為參考點(diǎn)的個數(shù)，組數(shù)為注意力頭的數(shù)量。

（3）計(jì)算注意力權(quán)重矩陣  的方法為  過一個nn.Linear和一個F.softmax，得到每個頭的注意力權(quán)重。

如圖2所示，分頭計(jì)算完的注意力最終會拼接到一起，然后最后過一個nn.Linear得到輸入  的最終輸出。

圖3. Multi-Scale Feature Maps

多尺度的Deformable Attention模塊也是在多尺度特征圖上計(jì)算的。多尺度的特征融合方法則是取了骨干網(wǎng)(ResNet)最后三層的特征圖C3，C4，C5，并且用了一個Conv3x3 Stride2的卷積得到了一個C6構(gòu)成了四層特征圖。特別的是會通過卷積操作將通道數(shù)量統(tǒng)一為256(也就是token的數(shù)量)，然后在這四個特征圖上運(yùn)行Deformable Attention Module并且進(jìn)行直接相加得到最終輸出。其中Deformable Attention Module算子的pytorch實(shí)現(xiàn)如下：

def ms_deform_attn_core_pytorch(value, value_spatial_shapes, sampling_locations, attention_weights):
    # for debug and test only,
    # need to use cuda version instead
    N_, S_, M_, D_ = value.shape # batch size, number token, number head, head dims
    # Lq_: number query, L_: level number, P_: sampling number采樣點(diǎn)數(shù)
    _, Lq_, M_, L_, P_, _ = sampling_locations.shape
    # 按照level劃分value
    value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1)
    # [0, 1] -> [-1, 1] 因?yàn)橐獫M足F.grid_sample的輸入要求
    sampling_grids = 2 * sampling_locations - 1
    sampling_value_list = []
    for lid_, (H_, W_) in enumerate(value_spatial_shapes):
        # N_, H_*W_, M_, D_ -> N_, H_*W_, M_*D_ -> N_, M_*D_, H_*W_ -> N_*M_, D_, H_, W_
        value_l_ = value_list[lid_].flatten(2).transpose(1, 2).reshape(N_*M_, D_, H_, W_)
        # N_, Lq_, M_, P_, 2 -> N_, M_, Lq_, P_, 2 -> N_*M_, Lq_, P_, 2
        sampling_grid_l_ = sampling_grids[:, :, :, lid_].transpose(1, 2).flatten(0, 1)
        # N_*M_, D_, Lq_, P_
        # 用雙線性插值從feature map上獲取value，因?yàn)閙ask的原因越界所以要zeros的方法進(jìn)行填充
        sampling_value_l_ = F.grid_sample(value_l_, sampling_grid_l_,
                                          mode='bilinear', padding_mode='zeros', align_corners=False)
        sampling_value_list.append(sampling_value_l_)
    # (N_, Lq_, M_, L_, P_) -> (N_, M_, Lq_, L_, P_) -> (N_, M_, 1, Lq_, L_*P_)
    attention_weights = attention_weights.transpose(1, 2).reshape(N_*M_, 1, Lq_, L_*P_)
    # 不同scale計(jì)算出的multi head attention 進(jìn)行相加，返回output后還需要過一個Linear層
    output = (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights).sum(-1).view(N_, M_*D_, Lq_)
    return output.transpose(1, 2).contiguous()

完整的Multi-Scale Deformable Attention模塊代碼如下：

class MSDeformAttn(nn.Module):
    def __init__(self, d_model=256, n_levels=4, n_heads=8, n_points=4):
        """
        Multi-Scale Deformable Attention Module
        :param d_model      hidden dimension
        :param n_levels     number of feature levels
        :param n_heads      number of attention heads
        :param n_points     number of sampling points per attention head per feature level
        """
        super().__init__()
        if d_model % n_heads != 0:
            raise ValueError('d_model must be divisible by n_heads, but got {} and {}'.format(d_model, n_heads))
        _d_per_head = d_model // n_heads
        # you'd better set _d_per_head to a power of 2 which is more efficient in our CUDA implementation
        if not _is_power_of_2(_d_per_head):
            warnings.warn("You'd better set d_model in MSDeformAttn to make the dimension of each attention head a power of 2 "
                          "which is more efficient in our CUDA implementation.")

        self.im2col_step = 64

        self.d_model = d_model
        self.n_levels = n_levels
        self.n_heads = n_heads
        self.n_points = n_points

        self.sampling_offsets = nn.Linear(d_model, n_heads * n_levels * n_points * 2)
        self.attention_weights = nn.Linear(d_model, n_heads * n_levels * n_points)
        self.value_proj = nn.Linear(d_model, d_model)
        self.output_proj = nn.Linear(d_model, d_model)

        self._reset_parameters()

    def _reset_parameters(self):
        constant_(self.sampling_offsets.weight.data, 0.)
        thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads)
        grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
        grid_init = (grid_init / grid_init.abs().max(-1, keepdim=True)[0]).view(self.n_heads, 1, 1, 2).repeat(1, self.n_levels, self.n_points, 1)
        for i in range(self.n_points):
            grid_init[:, :, i, :] *= i + 1
        with torch.no_grad():
            self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))
        constant_(self.attention_weights.weight.data, 0.)
        constant_(self.attention_weights.bias.data, 0.)
        xavier_uniform_(self.value_proj.weight.data)
        constant_(self.value_proj.bias.data, 0.)
        xavier_uniform_(self.output_proj.weight.data)
        constant_(self.output_proj.bias.data, 0.)

    def forward(self, query, reference_points, input_flatten, input_spatial_shapes, input_level_start_index, input_padding_mask=None):
        """
        :param query                       (N, Length_{query}, C)
        :param reference_points            (N, Length_{query}, n_levels, 2), range in [0, 1], top-left (0,0), bottom-right (1, 1), including padding area
                                        or (N, Length_{query}, n_levels, 4), add additional (w, h) to form reference boxes
        :param input_flatten               (N, \sum_{l=0}^{L-1} H_l \cdot W_l, C)
        :param input_spatial_shapes        (n_levels, 2), [(H_0, W_0), (H_1, W_1), ..., (H_{L-1}, W_{L-1})]
        :param input_level_start_index     (n_levels, ), [0, H_0*W_0, H_0*W_0+H_1*W_1, H_0*W_0+H_1*W_1+H_2*W_2, ..., H_0*W_0+H_1*W_1+...+H_{L-1}*W_{L-1}]
        :param input_padding_mask          (N, \sum_{l=0}^{L-1} H_l \cdot W_l), True for padding elements, False for non-padding elements

        :return output                     (N, Length_{query}, C)
        """
        # query是 src + positional encoding
        # input_flatten是src，沒有位置編碼
        N, Len_q, _ = query.shape
        N, Len_in, _ = input_flatten.shape
        assert (input_spatial_shapes[:, 0] * input_spatial_shapes[:, 1]).sum() == Len_in
        # 根據(jù)input_flatten得到v
        value = self.value_proj(input_flatten)
        if input_padding_mask is not None:
            value = value.masked_fill(input_padding_mask[..., None], float(0))
        # 多頭注意力 根據(jù)頭的個數(shù)將v等分
        value = value.view(N, Len_in, self.n_heads, self.d_model // self.n_heads)
        # 根據(jù)query得到offset偏移量和attention weights注意力權(quán)重
        sampling_offsets = self.sampling_offsets(query).view(N, Len_q, self.n_heads, self.n_levels, self.n_points, 2)
        attention_weights = self.attention_weights(query).view(N, Len_q, self.n_heads, self.n_levels * self.n_points)
        attention_weights = F.softmax(attention_weights, -1).view(N, Len_q, self.n_heads, self.n_levels, self.n_points)
        # N, Len_q, n_heads, n_levels, n_points, 2
        if reference_points.shape[-1] == 2:
            offset_normalizer = torch.stack([input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1)
            sampling_locations = reference_points[:, :, None, :, None, :] \
                                 + sampling_offsets / offset_normalizer[None, None, None, :, None, :]
        elif reference_points.shape[-1] == 4:
            sampling_locations = reference_points[:, :, None, :, None, :2] \
                                 + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5
        else:
            raise ValueError(
                'Last dim of reference_points must be 2 or 4, but get {} instead.'.format(reference_points.shape[-1]))
        output = MSDeformAttnFunction.apply(
            value, input_spatial_shapes, input_level_start_index, sampling_locations, attention_weights, self.im2col_step)
        output = self.output_proj(output)
        return output

詳細(xì)代碼注釋如下，iterative bounding box refinement和two stage改進(jìn)方法的Encoder不變。

class DeformableTransformerEncoderLayer(nn.Module):
    def __init__(self,
                 d_model=256, d_ffn=1024,
                 dropout=0.1, activation="relu",
                 n_levels=4, n_heads=8, n_points=4):
        super().__init__()

        # self attention
        self.self_attn = MSDeformAttn(d_model, n_levels, n_heads, n_points)
        self.dropout1 = nn.Dropout(dropout)
        self.norm1 = nn.LayerNorm(d_model)

        # ffn
        self.linear1 = nn.Linear(d_model, d_ffn)
        self.activation = _get_activation_fn(activation)
        self.dropout2 = nn.Dropout(dropout)
        self.linear2 = nn.Linear(d_ffn, d_model)
        self.dropout3 = nn.Dropout(dropout)
        self.norm2 = nn.LayerNorm(d_model)

    @staticmethod
    def with_pos_embed(tensor, pos):
        return tensor if pos is None else tensor + pos

    def forward_ffn(self, src):
        src2 = self.linear2(self.dropout2(self.activation(self.linear1(src))))
        src = src + self.dropout3(src2)
        src = self.norm2(src)
        return src

    def forward(self, src, pos, reference_points, spatial_shapes, level_start_index, padding_mask=None):
        # self attention
        src2 = self.self_attn(self.with_pos_embed(src, pos), reference_points, src, spatial_shapes, level_start_index, padding_mask)
        src = src + self.dropout1(src2)
        src = self.norm1(src)

        # ffn
        src = self.forward_ffn(src)

        return src

class DeformableTransformerEncoder(nn.Module):
    def __init__(self, encoder_layer, num_layers):
        super().__init__()
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers

    @staticmethod
    def get_reference_points(spatial_shapes, valid_ratios, device):
        reference_points_list = []
        for lvl, (H_, W_) in enumerate(spatial_shapes):
            # 從0.5到H-0.5采樣H個點(diǎn)，W同理 這個操作的目的也就是為了特征圖的對齊
            ref_y, ref_x = torch.meshgrid(torch.linspace(0.5, H_ - 0.5, H_, dtype=torch.float32, device=device),
                                          torch.linspace(0.5, W_ - 0.5, W_, dtype=torch.float32, device=device))
            ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, lvl, 1] * H_)
            ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, lvl, 0] * W_)
            ref = torch.stack((ref_x, ref_y), -1)
            reference_points_list.append(ref)
        reference_points = torch.cat(reference_points_list, 1)
        reference_points = reference_points[:, :, None] * valid_ratios[:, None]
        return reference_points

    def forward(self, src, spatial_shapes, level_start_index, valid_ratios, pos=None, padding_mask=None):
        output = src
        reference_points = self.get_reference_points(spatial_shapes, valid_ratios, device=src.device)
        for _, layer in enumerate(self.layers):
            output = layer(output, pos, reference_points, spatial_shapes, level_start_index, padding_mask)

        return output