【傳知代碼】DETR[端到端目標檢測]（論文復現）

前言：想象一下，當自動駕駛汽車行駛在繁忙的街道上，DETR能夠實時識別出道路上的行人、車輛、交通標志等目標，并準確預測出它們的位置和軌跡。這對于提高自動駕駛的安全性、減少交通事故具有重要意義。同樣，在安防監控、醫療影像分析等領域，DETR也展現出了巨大的應用潛力，如今，一項名為DETR（Detection Transformer）的創新技術，猶如一股清流，為這一領域帶來了革命性的變革。DETR，這個聽起來有些神秘而高深的名詞，實際上是一種基于Transformer架構的端到端目標檢測模型。它摒棄了傳統方法中繁瑣的錨框和候選區域生成步驟，直接通過Transformer的強大能力，將圖像中的目標信息與上下文信息相融合，實現了對目標的精準定位和分類。

本文所涉及所有資源均在傳知代碼平臺可獲取

概述

演示效果

核心代碼

寫在最后

概述

????????在進行目標檢測時，需要大量手動設計的組件，比如非極大值抑制（NMS）和基于人工經驗生成的先驗框（Anchor）等。DETR在其文章中，將目標檢測視為一個直接的集合預測任務，從而減少了對人工組件設計的依賴，并使目標檢測流程更為簡潔。當提供一組固定的、可學習的目標查詢DETR來推斷目標與全局圖像之間的上下文關系時，由于DETR沒有先驗框的限制，這將使其在預測較大物體時表現得更為出色。

????????如下圖展示的是DETR的核心框架。由于直接使用了transformer的結構，這導致模型的計算需求增加。因此，DETR首先利用CNN卷積神經網絡來提取特征，這種方法生成的特征圖通常會降低32倍的采樣。接下來，我們將提取出的特征圖傳輸到Transformer的encoder結構中，以實現自注意力的交互，從而揭示特征圖中每一個像素與其他像素的相互關系。decoder首先為用戶預設了N個查詢。這些查詢首先通過自注意力機制去除模型中的多余框，然后與來自Encoder的特征交互，生成數量為N的查詢。這些查詢通過線性層生成模型預測的類別和相應的邊界框輸出，最終完成預測：

實驗中N個數據比一幅圖包含全部對象更多，計算損失函數時DETR先用匈牙利算法找到合適匹配方式。然后去算bbox及分類損失值。鑒于L1L1損失函數對不同尺寸的邊界框產生的誤差存在差異，我們決定使用GIoUGIoU損失函數來補償這些誤差。如下圖，為DETR更為詳盡的圖示：

主干網絡方面：

????????針對于一張通道數大小為3的圖片，首先經過CNN的骨干網絡，得到一個通道數為2048（這個數據由我們手動設定），長寬分別為原始圖像大小132321?的特征圖f∈RC×H×Wf∈RC×H×。

Transformer編碼器：

????????首先，通過1×11×1的卷積方法，我們將特征圖ff的通道維數從CC減少到了更低的dd維度，并據此生成了一個新的特征圖z0∈Rd×H×Wz0∈Rd×H×W，編碼器希望有序列做輸入，所以我們把z0z0?個空間維度壓縮成1維，生成d×HWd×HW特征圖。
每一個編碼器層都配備了一個統一的架構，該架構由一個多頭自注意力模塊和一個前饋網絡（FFN）共同構成。由于Transformer架構具有置換不變性（對輸入序列進行排序更改而不會對輸出結果進行更改），我們用維度大小相同的位置編碼來彌補這個缺點，位置編碼被添加到每個注意力層的輸入中。
Transformer解碼器：

????????DETR與標準Transformer架構中的decoder有所不同，因為它并未使用掩碼技術，這意味著N個預測的邊界框可以被同時輸出。

鑒于解碼器依然保持置換不變性，我們選擇了可學習的位置編碼作為其輸入嵌入方式，并將其命名為object query。這種object query經由若干層結構最后被轉換到輸出邊界框上并經由FFN結構產生N個坐標點以及分類后之物體。

????????下圖所示是模型Transformer的主要結構，來自CNN主干網絡的圖像特征被送到transformer編碼器中，在每個多頭自注意力機制中與空間位置編碼相加作為多頭自注意力機制的鍵和查詢，（生成q,k,v需要矩陣相乘，并不是一個直接的結果）。作為在解碼器和編碼器進行注意力機制計算之前，首先object query需要進行一個自注意力機制，該步驟是為了去除模型中的冗余框：

演示效果

使用一個GPU進行模型訓練、驗證和可視化，命令如下：

# 模型訓練
python -m torch.distributed.launch --nproc_per_node=1 --use_env main.py --coco_path data/coco # 模型驗證
python main.py --batch_size 2 --no_aux_loss --eval --resume ckpt/detr-r50-e632da11.pth --coco_path data/coco# 模型可視化
python imshow.py

部署項目方式如下：

# 首先安裝相應版本的PyTorch 1.5+和torchvision 0.6+ ，如果有GPU則安裝GPU版本的，沒有安裝相應cpu版本的，注意linux和window之間的區別
conda install -c pytorch pytorch torchvision
# 安裝pycococtools(在COCO數據集上進行預測)和scipy(為了訓練)
conda install cython scipy
pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

從http://cocodataset.org下載COCO2017的train和val圖像，相應地annotation,具體如下截圖所示：

將數據集按照下面的形式進行擺放：

data/coco/annotations/  # annotation json filestrain2017/    # train imagesval2017/      # val images

從detr-r50-e632da11.pth下載相應的權重，并命名為ckpt/detr-r50-e632da11.pth，放在ckpt文件夾下，如下圖所示：

使用DETR進行目標檢測，效果如下：?

使用DETR交叉注意力機制可視化如下：

DETR自注意力機制可視化，query表示當前物體的標號，下方對應的是相應的名稱，下方顯示的點可以人工手動調整：

核心代碼

下面這段代碼實現了一個目標檢測模型 DETR（DEtection TRansformer），它使用了 Transformer 架構進行目標檢測，在 __init__ 函數中，模型接受了一個 backbone 模型、一個 transformer 模型、目標類別數 num_classes、最大檢測框個數 num_queries 和一個參數 aux_loss。其中，backbone 模型用于提取特征，transformer 模型用于處理特征和進行目標檢測。模型的輸出包括分類 logits 和檢測框坐標，以及可選的輔助損失。

在 forward 函數中，模型接受了一個 NestedTensor，其中 samples.tensor 是一個批次的圖像，samples.mask 是一個二進制掩碼，表示每個圖像中的有效像素。首先，模型使用 backbone 模型提取特征和位置編碼。然后，模型使用 transformer 模型對特征和位置編碼進行處理，得到分類 logits 和檢測框坐標。最后，模型將分類 logits 和檢測框坐標輸出為字典，其中 pred_logits 表示分類 logits，pred_boxes 表示檢測框坐標。

在 _set_aux_loss 函數中，模型處理輔助損失。這里使用了一個 workaround，將輸出的字典轉換為一個列表，每個元素包含分類 logits 和檢測框坐標。這樣做是為了讓 torchscript 能夠正常工作，因為它不支持非同構值的字典。

class DETR(nn.Module):""" This is the DETR module that performs object detection """def __init__(self, backbone, transformer, num_classes, num_queries, aux_loss=False):""" Initializes the model.Parameters:backbone: torch module of the backbone to be used. See backbone.pytransformer: torch module of the transformer architecture. See transformer.pynum_classes: number of object classesnum_queries: number of object queries, ie detection slot. This is the maximal number of objectsDETR can detect in a single image. For COCO, we recommend 100 queries.aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used."""super().__init__()self.num_queries = num_queriesself.transformer = transformerhidden_dim = transformer.d_modelself.class_embed = nn.Linear(hidden_dim, num_classes + 1)self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)self.query_embed = nn.Embedding(num_queries, hidden_dim)self.input_proj = nn.Conv2d(backbone.num_channels, hidden_dim, kernel_size=1)self.backbone = backboneself.aux_loss = aux_lossdef forward(self, samples: NestedTensor):""" The forward expects a NestedTensor, which consists of:- samples.tensor: batched images, of shape [batch_size x 3 x H x W]- samples.mask: a binary mask of shape [batch_size x H x W], containing 1 on padded pixelsIt returns a dict with the following elements:- "pred_logits": the classification logits (including no-object) for all queries.Shape= [batch_size x num_queries x (num_classes + 1)]- "pred_boxes": The normalized boxes coordinates for all queries, represented as(center_x, center_y, height, width). These values are normalized in [0, 1],relative to the size of each individual image (disregarding possible padding).See PostProcess for information on how to retrieve the unnormalized bounding box.- "aux_outputs": Optional, only returned when auxilary losses are activated. It is a list ofdictionnaries containing the two above keys for each decoder layer."""if isinstance(samples, (list, torch.Tensor)):samples = nested_tensor_from_tensor_list(samples)# backbone 網絡進行了兩個操作，分別是獲取特征圖和位置編碼features, pos = self.backbone(samples)src, mask = features[-1].decompose()assert mask is not None# input_proj: src: [2,2048,28,38]->[2,256,28,38] 改變特征圖的通道維數# mask: [2,28,38] mask的通道維數為1 pos: [2,256,28,38] query表示查詢，也就是圖片里面可能有多少物體的個數hs = self.transformer(self.input_proj(src), mask, self.query_embed.weight, pos[-1])[0]outputs_class = self.class_embed(hs)outputs_coord = self.bbox_embed(hs).sigmoid()# 都只使用最后一層decoder輸出的結果out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]}if self.aux_loss:out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord)return out@torch.jit.unuseddef _set_aux_loss(self, outputs_class, outputs_coord):# this is a workaround to make torchscript happy, as torchscript# doesn't support dictionary with non-homogeneous values, such# as a dict having both a Tensor and a list.return [{'pred_logits': a, 'pred_boxes': b}for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]

下面這段代碼實現了一個 Transformer 模型，用于對輸入特征進行編碼和解碼，首先，模型將輸入特征和位置編碼展平，并進行轉置，得到形狀為 [HW, N, C] 的張量。然后，模型將查詢編碼重復 N 次，并將掩碼展平，以便在解碼器中使用。接下來，模型使用編碼器對輸入特征進行編碼，并使用解碼器對編碼后的特征進行解碼。最后，模型將解碼器的輸出進行轉置，得到形狀為 [batch_size, num_queries, d_model] 的張量，并將編碼器的輸出進行轉置和重構，得到與輸入特征相同的形狀，如下：

class Transformer(nn.Module):def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,activation="relu", normalize_before=False,return_intermediate_dec=False):super().__init__()encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward,dropout, activation, normalize_before)encoder_norm = nn.LayerNorm(d_model) if normalize_before else Noneself.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward,dropout, activation, normalize_before)decoder_norm = nn.LayerNorm(d_model)self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm,return_intermediate=return_intermediate_dec)self._reset_parameters()self.d_model = d_modelself.nhead = nheaddef _reset_parameters(self):for p in self.parameters():if p.dim() > 1:nn.init.xavier_uniform_(p)def forward(self, src, mask, query_embed, pos_embed):# flatten NxCxHxW to HWxNxC [2,256,28,38]bs, c, h, w = src.shape# src: [2,256,28,38]->[2,256,28*38]->[1064,2,256]# pos_embed: [2,256,28,38]->[2,256,28*38]->[1064,2,256]src = src.flatten(2).permute(2, 0, 1)pos_embed = pos_embed.flatten(2).permute(2, 0, 1)# query_embed:[100,256]->[100,1,256]->[100,2,256]query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1)# mask: [2,28,38]->[2,1064]mask = mask.flatten(1)# 其實也是一個位置編碼，表示目標的信息，一開始被初始化為0 [100,2,256]tgt = torch.zeros_like(query_embed)# memory的shape和src的一樣是[1064,2,256]memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)hs = self.decoder(tgt, memory, memory_key_padding_mask=mask,pos=pos_embed, query_pos=query_embed)# hs 不止輸出最后一層的結構，而是輸出解碼器所有層結構的輸出情況# hs: [6,100,2,256]->[6,2,100,256] [depth,batch_size,num_query,channel]# 一般只使用最后一層特征所以未hs[-1]return hs.transpose(1, 2), memory.permute(1, 2, 0).view(bs, c, h, w)

寫在最后

????????DETR以其獨特的視角和創新的架構，徹底改變了目標檢測的傳統流程。它摒棄了復雜的預處理步驟，如錨框生成和非極大值抑制，轉而采用了一種簡潔而高效的設計。通過Transformer的自注意力機制，DETR能夠捕捉圖像中各個部分之間的長距離依賴關系，從而更準確地預測目標的位置和類別。

????????DETR的成功并非偶然。它基于Transformer的強大能力，將圖像特征提取、目標定位和分類任務全部整合在一個模型中，實現了真正的端到端訓練。這種設計不僅簡化了檢測過程，還提高了模型的整體優化效果。更重要的是，DETR的“集合預測”機制允許模型以并行的方式預測所有目標，無需繁瑣的排序或篩選操作，進一步提升了檢測效率。

詳細復現過程的項目源碼、數據和預訓練好的模型可從該文章下方附件獲取。