基於 EasyCV 復現 DETR 和 DAB-DETR，Object Query 的正確打開方式
DETR 是最近幾年最新的目標檢測框架，第一個真正意義上的端到端檢測演算法，省去了繁瑣的 RPN、anchor 和 NMS 等操作，直接輸入圖片輸出檢測框。DETR 的成功主要歸功於 Transformer 強大的建模能力，還有匈牙利匹配演算法解決了如何通過學習的方式 one-to-one 的匹配檢測框和目標框。
雖然 DETR 可以達到跟 Mask R-CNN 相當的精度，但是訓練 500 個 epoch、收斂速度慢，小目標精度低的問題都飽受詬病。後續一系列的工作都圍繞著這幾個問題展開，其中最精彩的要屬 Deformable DETR，也是如今檢測的刷榜必備，Deformable DETR 的貢獻不單單只是將 Deformable Conv 推廣到了 Transformer 上，更重要的是提供了很多訓練好 DETR 檢測框架的技巧，比如模仿 Mask R-CNN 框架的 two-stage 做法，如何將 query embed 拆分成 content 和 reference points 兩部分組成，如何將 DETR 拓展到多尺度訓練，還有通過 look forward once 進行 boxes 預測等技巧，在 Deformable DETR 之後，大家似乎找到了如何打開 DETR 框架的正確方式。其中對 object query 代表什麼含義，以及如何更好的利用 object query 做檢測，產生了許多有價值的工作，比如 Anchor DETR、Conditional DETR 等等，其中 DAB-DETR 做的尤為徹底。DAB-DETR 將 object query 看成是 content 和 reference points 兩個部分，其中 reference points 顯示的表示成 xywh 四維向量，然後通過 decoder 預測 xywh 的殘差對檢測框迭代更新，另外還通過 xywh 向量引入位置注意力，幫助 DETR 加快收斂速度，本文將基於 EasyCV 復現的 DETR 和 DAB-DETR 演算法詳細介紹一下如何正確的使用 object query 來提升 DETR 檢測框架的性能。
DETR
DETR 使用 set loss function 作為監督訊號來進行端到端訓練，然後同時預測所有目標，其中 set loss function 使用 bipartite matching 演算法將 pred 目標和 gt 目標匹配起來。直接將目標檢測任務看成 set prediction 問題，使訓練過程變的簡潔，並且避免了 anchor、NMS 等複雜處理。
DETR 主要貢獻有兩個部分：architecture 和 set prediction loss。
1.Architecture
DETR 先用 CNN 將輸入影像 embedding 成一個二維表徵，然後將二維表徵轉換成一維表徵並結合 positional encoding 一起送入 encoder，decoder 將少量固定數量的已學習的 object queries(可以理解為 positional embeddings)和 encoder 的輸出作為輸入。最後將 decoder 得到的每個 output embdding 傳遞到一個共享的前饋網路(FFN)，該網路可以預測一個檢測結果(包括類和邊框)或著「沒有目標」的類。
1.1 Transformer
1.1.1 Encoder
將 Backbone 輸出的 feature map 轉換成一維表徵，得到 特徵圖，然後結合 positional encoding 作為 Encoder 的輸入。每個 Encoder 都由 Multi-Head Self-Attention 和 FFN 組成。和 Transformer Encoder 不同的是，因為 Encoder 具有位置不變性，DETR 將 positional encoding 添加到每一個 Multi-Head Self-Attention 中，來保證目標檢測的位置敏感性。
# 一層encoder程式碼如下 class TransformerEncoderLayer(nn.Module): def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='relu', normalize_before=False): super().__init__() self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) # Implementation of Feedforward model self.linear1 = nn.Linear(d_model, dim_feedforward) self.dropout = nn.Dropout(dropout) self.linear2 = nn.Linear(dim_feedforward, d_model) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.activation = _get_activation_fn(activation) self.normalize_before = normalize_before def with_pos_embed(self, tensor, pos: Optional[Tensor]): return tensor if pos is None else tensor + pos def forward(self, src, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None): q = k = self.with_pos_embed(src, pos) src2 = self.self_attn( q, k, value=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0] src = src + self.dropout1(src2) src = self.norm1(src) src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) src = src + self.dropout2(src2) src = self.norm2(src) return src
1.1.2 Decoder
因為 Decoder 也具有位置不變性，Decoder 的N個 object query(可以理解為學習不同 object 的 positional embedding)必須是不同，以便生成不同 object 的 embedding，並且同時把它們添加到每一個 Multi-Head Attention 中。N個 object queries 通過 Decoder 轉換成一個 output embedding，然後 output embedding 通過 FFN 獨立解碼出N個預測結果，包含 box 和 class。對輸入 embedding 同時使用 Self-Attention 和 Encoder-Decoder Attention，模型可以利用目標的相互關係來進行全局推理。和 Transformer Decoder 不同的是，DETR 的每個 Decoder 並行輸出N個對象，Transformer Decoder 使用的是自回歸模型，串列輸出N個對象，每次只能預測一個輸出序列的一個元素。
# 一層decoder程式碼如下 class TransformerDecoderLayer(nn.Module): def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='relu', normalize_before=False): super().__init__() self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) self.multihead_attn = nn.MultiheadAttention( d_model, nhead, dropout=dropout) # Implementation of Feedforward model self.linear1 = nn.Linear(d_model, dim_feedforward) self.dropout = nn.Dropout(dropout) self.linear2 = nn.Linear(dim_feedforward, d_model) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.norm3 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.dropout3 = nn.Dropout(dropout) self.activation = _get_activation_fn(activation) self.normalize_before = normalize_before def with_pos_embed(self, tensor, pos: Optional[Tensor]): return tensor if pos is None else tensor + pos def forward(self, tgt, memory, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None, query_pos: Optional[Tensor] = None): q = k = self.with_pos_embed(tgt, query_pos) tgt2 = self.self_attn( q, k, value=tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0] tgt = tgt + self.dropout1(tgt2) tgt = self.norm1(tgt) tgt2 = self.multihead_attn( query=self.with_pos_embed(tgt, query_pos), key=self.with_pos_embed(memory, pos), value=memory, attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask)[0] tgt = tgt + self.dropout2(tgt2) tgt = self.norm2(tgt) tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) tgt = tgt + self.dropout3(tgt2) tgt = self.norm3(tgt) return tgt
1.1.3 FFNFFN
由 3 層 perceptron 和一層 linear projection 組成。FFN 預測出 box 的歸一化中心坐標、長、寬和 class。DETR 預測的是固定數量的N個 box 的集合，並且N通常比實際目標數要大的(其中 DETR 默認設置為 100 個，而 DAB-DETR 設置為 300 個)，並且使用一個額外的空類來表示預測得到的 box 不存在目標。
class MLP(nn.Module): """ Very simple multi-layer perceptron (also called FFN)""" def __init__(self, input_dim, hidden_dim, output_dim, num_layers): super().__init__() self.num_layers = num_layers h = [hidden_dim] * (num_layers - 1) self.layers = nn.ModuleList( nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])) def forward(self, x): for i, layer in enumerate(self.layers): x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x) return x
2.Set prediction loss
DETR 模型訓練的主要困難是如何根據 gt 衡量預測結果(類別、位置、數量)。DETR 提出的 loss 函數可以產生 pred 和 gt 的最優雙邊匹配(確定 pred 和 gt 的一對一關係)，然後優化 loss。將y表示為 gt 的集合， 表示為N個預測結果的集合。假設N大於圖片目標數，y可以認為是用空類(無目標)填充的大小為N的集合。搜索兩個集合N個元素
的不同排列順序，使得 loss 儘可能的小的排列順序即為二分圖最大匹配(Bipartite Matching)，公式如下：
其中
表示 pred 和 gt 關於σ(i)元素i的匹配 loss。其中二分圖匹配通過匈牙利演算法(Hungarian algorithm)得到。匹配 loss 同時考慮了 pred class 和 pred box 的準確性。每個 gt 的元素i可以看成yi=(ci,bi)，ci表示 class label(可能是空類)bi表示 gt box，將元素i二分圖匹配指定的 pred class 表示為
，
pred box 表示為
。
第一步先找到一對一匹配的 pred 和 gt，第二步再計算 hungarian loss。hungarian loss 公式如下：
其中 結合了 L1 loss 和 generalized IoU loss，公式如下：
# HungarianMatcher通過計算出cost_bbox,cost_class,cost_giou來一對一匹配預測框和gt框，然後返回匹配的索引對，最後通過索引對計算出loss值 # Final cost matrix C = self.cost_bbox * cost_bbox + self.cost_class * cost_class + self.cost_giou * cost_giou C = C.view(bs, num_queries, -1).cpu() sizes = [len(v['boxes']) for v in targets] indices = [ linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1)) ] return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
DAB-DETR
DAB-DETR 將 object query 看成是 content 和 reference points 兩個部分，其中 reference points 顯示的表示成 xywh 四維向量，然後通過 decoder 預測 xywh 的殘差對檢測框迭代更新，另外還通過 xywh 向量引入位置注意力，幫助 DETR 加快收斂速度。
在 DAB-DETR 之前，有許多工作對如何設置 reference points 進行過深入的探索：Conditional DETR 通過 256 維的可學習向量學習得到 xy 參考點，然後將位置資訊引入 transformer decoder 中；Anchor DETR 參考點看成是 xy，然後通過學習的方式得到 256 維的向量，將位置資訊引入 transformer decoder 中，並且通過逐級迭代得到檢測框的 xy；Defomable DETR 則是通過 256 維可學習向量得到 xywh 參考 anchor，通過逐級迭代得到檢測框；DAB-DETR 則更為徹底，吸百家之長，通過 xywh 學習 256 維的向量，將位置資訊引入 transformer decoder 中，並且通過逐級迭代得到檢測框。至此，reference points 的使用方式逐漸明朗起來，顯示的表示為 xywh，然後學習成 256 維向量，引入位置資訊，每層 transformer decoder 學習 xywh 的殘差，逐級疊加得到最後的檢測框。
# DAB-DETR將object query顯示的拆分為content和pos兩種屬性 # 將query_embed顯示的表示為xywh，表示pos屬性，通過MLP學習成256維的pos特徵 self.query_embed = nn.Embedding(num_queries, query_dim) # get sine embedding for the query vector reference_points = self.query_embed.sigmoid() obj_center = reference_points[..., :2] query_sine_embed = gen_sineembed_for_position(obj_center) query_pos = self.ref_point_head(query_sine_embed) # content_embed初始化為全0的256維特徵 tgt = torch.zeros( self.num_queries, bs, self.embed_dims, device=query_embed.device)
另外，DAB-DETR 為了更充分的利用 xywh 這種更為顯示的 reference points 表示方式，進一步的引入了 Width & Height-Modulated Multi-Head Cross-Attention，其實簡單來講就是在 cross-attention 中引入位置 xywh 得到的位置注意力，這一點改進可以極大的加快 decoder 的收斂速度，因為原始的 DETR 相當於是在全圖學習到位置注意力，DAB-DETR 可以直接關注到關鍵位置，這也是 Deformable DETR 為啥能加快收斂的原因，本質就是更關鍵的稀疏位置取樣可以加快 decoder 收斂速度。
# 通過MLP的學習，調整query_sine_embed的attn位置，進一步加快收斂速度 # modulated HW attentions if self.modulate_hw_attn: refHW_cond = self.ref_anchor_head( output).sigmoid() # nq, bs, 2 query_sine_embed[..., self.d_model // 2:] *= (refHW_cond[..., 0] / obj_center[..., 2]).unsqueeze(-1) query_sine_embed[..., :self.d_model // 2] *= (refHW_cond[..., 1] / obj_center[..., 3]).unsqueeze(-1)
復現結果
Tutorial
接下來，我們將通過一個實際的例子介紹如何基於 EasyCV 進行 DAB-DETR 演算法的訓練，也可以在該鏈接查看詳細步驟。
一、安裝依賴包
如果是在本地開發環境運行，可以參考該鏈接安裝環境。若使用 PAI-DSW 進行實驗則無需安裝相關依賴，在 PAI-DSW docker 中已內置相關環境。二、數據準備
你可以下載COCO2017數據，也可以使用我們提供了示例 COCO 數據
wget //pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/data/small_coco_demo/small_coco_demo.tar.gz && tar -zxf small_coco_demo.tar.gz mkdir -p data/ && mv small_coco_demo data/coco
data/coco 格式如下：
data/coco/ ├── annotations │ ├── instances_train2017.json │ └── instances_val2017.json ├── train2017 │ ├── 000000005802.jpg │ ├── 000000060623.jpg │ ├── 000000086408.jpg │ ├── 000000118113.jpg │ ├── 000000184613.jpg │ ├── 000000193271.jpg │ ├── 000000222564.jpg │ ... │ └── 000000574769.jpg └── val2017 ├── 000000006818.jpg ├── 000000017627.jpg ├── 000000037777.jpg ├── 000000087038.jpg ├── 000000174482.jpg ├── 000000181666.jpg ├── 000000184791.jpg ├── 000000252219.jpg ... └── 000000522713.jpg
二、模型訓練和評估
以 vitdet-base 為示例。在 EasyCV 中，使用配置文件的形式來實現對模型參數、數據輸入及增廣方式、訓練策略的配置，僅通過修改配置文件中的參數設置，就可以完成實驗配置進行訓練。可以直接下載示例配置文件。
查看 easycv 安裝位置
# 查看easycv安裝位置
import easycv
print(easycv.__file__)
export PythonPATH=$PYTHONPATH:root/EasyCV
執行訓練命令
單機8卡： CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m torch.distributed.launch -- nproc_per_node=8 -- master_port=29500 tools/train.py configs/detection/dab- detr/dab_detr_r50_8x2_50e_coco.p y --work_dir easycv/dab_detr - -launcher pytorch
執行評估命令
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5 ,6,7 python -m torch.distributed.launch -- nproc_per_node=8 -- master_port=29500 tools/eval.py configs/detection/dab- detr/dab_detr_r50_8x2_50e_coco.p y easycv/dab_detr/epoch_50.pth - -launcher pytorch --eval
Reference
程式碼實現：
DETR //github.com/alibaba/EasyCV/tree/master/easycv/models/detection/detectors/detr
DAB-DETR
//github.com/alibaba/EasyCV/tree/master/easycv/models/detection/detectors/dab_detr
EasyCV 往期分享
基於EasyCV復現ViTDet：單層特徵超越FPN//zhuanlan.zhihu.com/p/528733299
MAE 自監督演算法介紹和基於 EasyCV 的復現 //zhuanlan.zhihu.com/p/515859470
EasyCV 開源｜開箱即用的視覺自監督+Transformer 演算法庫 //zhuanlan.zhihu.com/p/505219993
