如何使用Transformer來(lái)做物體檢測(cè)？

導(dǎo)讀

本文為一個(gè)Facebook的目標(biāo)檢測(cè)Transformer （DETR）的完整指南，詳細(xì)介紹了DETR架構(gòu)的內(nèi)部工作方式以及代碼。

介紹

DEtection TRansformer （DETR）是Facebook研究團(tuán)隊(duì)巧妙地利用了Transformer 架構(gòu)開(kāi)發(fā)的一個(gè)目標(biāo)檢測(cè)模型。在這篇文章中，我將通過(guò)分析DETR架構(gòu)的內(nèi)部工作方式來(lái)幫助提供一些關(guān)于它的含義。下面，我將解釋一些結(jié)構(gòu)，但是如果你只是想了解如何使用模型，可以直接跳到代碼部分。

結(jié)構(gòu)

DETR模型由一個(gè)預(yù)訓(xùn)練的CNN骨干（如ResNet）組成，它產(chǎn)生一組低維特征集。這些特征被格式化為一個(gè)特征集合并添加位置編碼，輸入一個(gè)由Transformer組成的編碼器和解碼器中，和原始的Transformer論文中描述的Encoder-Decoder的使用方式非常的類似。解碼器的輸出然后被送入固定數(shù)量的預(yù)測(cè)頭，這些預(yù)測(cè)頭由預(yù)定義數(shù)量的前饋網(wǎng)絡(luò)組成。每個(gè)預(yù)測(cè)頭的輸出都包含一個(gè)類預(yù)測(cè)和一個(gè)預(yù)測(cè)框。損失是通過(guò)計(jì)算二分匹配損失來(lái)計(jì)算的。

該模型做出了預(yù)定義數(shù)量的預(yù)測(cè)，并且每個(gè)預(yù)測(cè)都是并行計(jì)算的。

CNN主干

假設(shè)我們的輸入圖像，有三個(gè)輸入通道。CNN backbone由一個(gè)（預(yù)訓(xùn)練過(guò)的）CNN（通常是ResNet）組成，我們用它來(lái)生成_C_個(gè)具有寬度W和高度H的低維特征（在實(shí)踐中，我們?cè)O(shè)置_C_=2048， W=W?/32和H=H?/32）。這留給我們的是C個(gè)二維特征，由于我們將把這些特征傳遞給一個(gè)transformer，每個(gè)特征必須允許編碼器將每個(gè)特征處理為一個(gè)序列的方式重新格式化。這是通過(guò)將特征矩陣扁平化為H?W向量，然后將每個(gè)向量連接起來(lái)來(lái)實(shí)現(xiàn)的。

扁平化的卷積特征再加上空間位置編碼，位置編碼既可以學(xué)習(xí)，也可以預(yù)定義。

The Transformer

Transformer幾乎與原始的編碼器-解碼器架構(gòu)完全相同。不同之處在于，每個(gè)解碼器層并行解碼N個(gè)（預(yù)定義的數(shù)目）目標(biāo)。該模型還學(xué)習(xí)了一組N個(gè)目標(biāo)的查詢，這些查詢是（類似于編碼器）學(xué)習(xí)出來(lái)的位置編碼。

目標(biāo)查詢

下圖描述了N=20個(gè)學(xué)習(xí)出來(lái)的目標(biāo)查詢（稱為prediction slots）如何聚焦于一張圖像的不同區(qū)域。

“我們觀察到，在不同的操作模式下，每個(gè)slot 都會(huì)學(xué)習(xí)特定的區(qū)域和框大小?！?—— DETR的作者

理解目標(biāo)查詢的直觀方法是想象每個(gè)目標(biāo)查詢都是一個(gè)人。每個(gè)人都可以通過(guò)注意力來(lái)查看圖像的某個(gè)區(qū)域。一個(gè)目標(biāo)查詢總是會(huì)問(wèn)圖像中心是什么，另一個(gè)總是會(huì)問(wèn)左下角是什么，以此類推。

使用PyTorch實(shí)現(xiàn)簡(jiǎn)單的DETR

import torchimport torch.nn as nnfrom torchvision.models import resnet50class SimpleDETR（nn.Module）：“”“Minimal Example of the Detection Transformer model with learned positional embedding”“” def __init__（self， num_classes， hidden_dim， num_heads， num_enc_layers， num_dec_layers）： super（SimpleDETR， self）.__init__（） self.num_classes = num_classes self.hidden_dim = hidden_dim self.num_heads = num_heads self.num_enc_layers = num_enc_layers self.num_dec_layers = num_dec_layers # CNN Backbone self.backbone = nn.Sequential（ *list（resnet50（pretrained=True）.children（））［：-2］） self.conv = nn.Conv2d（2048， hidden_dim， 1） # Transformer self.transformer = nn.Transformer（hidden_dim， num_heads， num_enc_layers， num_dec_layers） # Prediction Heads self.to_classes = nn.Linear（hidden_dim， num_classes+1） self.to_bbox = nn.Linear（hidden_dim， 4） # Positional Encodings self.object_query = nn.Parameter（torch.rand（100， hidden_dim）） self.row_embed = nn.Parameter（torch.rand（50， hidden_dim // 2） self.col_embed = nn.Parameter（torch.rand（50， hidden_dim // 2）） def forward（self， X）： X = self.backbone（X） h = self.conv（X） H， W = h.shape［-2：］ pos_enc = torch.cat（［ self.col_embed［：W］.unsqueeze（0）.repeat（H，1，1）， self.row_embed［：H］.unsqueeze（1）.repeat（1，W，1）］， dim=-1）.flatten（0，1）.unsqueeze（1） h = self.transformer（pos_enc + h.flatten（2）.permute（2，0，1）， self.object_query.unsqueeze（1）） class_pred = self.to_classes（h） bbox_pred = self.to_bbox（h）.sigmoid（） return class_pred， bbox_pred

二分匹配損失 （Optional）

讓為預(yù)測(cè)的集合，其中是包括了預(yù)測(cè)類別（可以是空類別）和包圍框的二元組，其中上劃線表示框的中心點(diǎn)， 和表示框的寬和高。設(shè)y為ground truth集合。假設(shè)y和_?_之間的損失為L(zhǎng)，每一個(gè)y?和_?_?之間的損失為L(zhǎng)?。由于我們是在集合的層次上工作，損失L必須是排列不變的，這意味著無(wú)論我們?nèi)绾闻判蝾A(yù)測(cè)，我們都將得到相同的損失。因此，我們想找到一個(gè)排列，它將預(yù)測(cè)的索引映射到ground truth目標(biāo)的索引上。在數(shù)學(xué)上，我們求解：

計(jì)算的過(guò)程稱為尋找最優(yōu)的二元匹配。這可以用匈牙利算法找到。但為了找到最優(yōu)匹配，我們需要實(shí)際定義一個(gè)損失函數(shù)，計(jì)算和之間的匹配成本。

回想一下，我們的預(yù)測(cè)包含一個(gè)邊界框和一個(gè)類?，F(xiàn)在讓我們假設(shè)類預(yù)測(cè)實(shí)際上是一個(gè)類集合上的概率分布。那么第_i_個(gè)預(yù)測(cè)的總損失將是類預(yù)測(cè)產(chǎn)生的損失和邊界框預(yù)測(cè)產(chǎn)生的損失之和。作者在http://arxiv.org/abs/1906.05909中將這種損失定義為邊界框損失和類預(yù)測(cè)概率的差異：

其中，是的argmax，是是來(lái)自包圍框的預(yù)測(cè)的損失，如果，則表示匹配損失為0。

框損失的計(jì)算為預(yù)測(cè)值與ground truth的L?損失和的GIOU損失的線性組合。同樣，如果你想象兩個(gè)不相交的框，那么框的錯(cuò)誤將不會(huì)提供任何有意義的上下文（我們可以從下面的框損失的定義中看到）。

其中，λ???和是超參數(shù)。注意，這個(gè)和也是面積和距離產(chǎn)生的誤差的組合。為什么會(huì)這樣呢？

可以把上面的等式看作是與預(yù)測(cè)相關(guān)聯(lián)的總損失，其中面積誤差的重要性是λ???，距離誤差的重要性是?，F(xiàn)在我們來(lái)定義GIOU損失函數(shù)。定義如下：

由于我們從已知的已知類的數(shù)目來(lái)預(yù)測(cè)類，那么類預(yù)測(cè)就是一個(gè)分類問(wèn)題，因此我們可以使用交叉熵?fù)p失來(lái)計(jì)算類預(yù)測(cè)誤差。我們將損失函數(shù)定義為每N個(gè)預(yù)測(cè)損失的總和：

為目標(biāo)檢測(cè)使用DETR

在這里，你可以學(xué)習(xí)如何加載預(yù)訓(xùn)練的DETR模型，以便使用PyTorch進(jìn)行目標(biāo)檢測(cè)。

加載模型

首先導(dǎo)入需要的模塊。

# Import required modulesimport torchfrom torchvision import transforms as T import requests # for loading images from webfrom PIL import Image # for viewing imagesimport matplotlib.pyplot as plt

下面的代碼用ResNet50作為CNN骨干從torch hub加載預(yù)訓(xùn)練的模型。其他主干請(qǐng)參見(jiàn)DETR github：https://github.com/facebookresearch/detr

detr = torch.hub.load（‘facebookresearch/detr’， ‘detr_resnet50’， pretrained=True）

加載一張圖像

要從web加載圖像，我們使用requests庫(kù)：

url = ‘https://www.tempetourism.com/wp-content/uploads/Postino-Downtown-Tempe-2.jpg’ # Sample imageimage = Image.open（requests.get（url， stream=True）.raw） plt.imshow（image）plt.show（）

設(shè)置目標(biāo)檢測(cè)的Pipeline

為了將圖像輸入到模型中，我們需要將PIL圖像轉(zhuǎn)換為張量，這是通過(guò)使用torchvision的transforms庫(kù)來(lái)完成的。

transform = T.Compose（［T.Resize（800）， T.ToTensor（）， T.Normalize（［0.485， 0.456， 0.406］， ［0.229， 0.224， 0.225］）］）

上面的變換調(diào)整了圖像的大小，將PIL圖像進(jìn)行轉(zhuǎn)換，并用均值-標(biāo)準(zhǔn)差對(duì)圖像進(jìn)行歸一化。其中［0.485，0.456，0.406］為各顏色通道的均值，［0.229，0.224，0.225］為各顏色通道的標(biāo)準(zhǔn)差。我們裝載的模型是預(yù)先在COCO Dataset上訓(xùn)練的，有91個(gè)類，還有一個(gè)表示空類（沒(méi)有目標(biāo)）的附加類。我們用下面的代碼手動(dòng)定義每個(gè)標(biāo)簽：

CLASSES = ［‘N/A’， ‘Person’， ‘Bicycle’， ‘Car’， ‘Motorcycle’， ‘Airplane’， ‘Bus’， ‘Train’， ‘Truck’， ‘Boat’， ‘Traffic-Light’， ‘Fire-Hydrant’， ‘N/A’， ‘Stop-Sign’， ‘Parking Meter’， ‘Bench’， ‘Bird’， ‘Cat’， ‘Dog’， ‘Horse’， ‘Sheep’， ‘Cow’， ‘Elephant’， ‘Bear’， ‘Zebra’， ‘Giraffe’， ‘N/A’， ‘Backpack’， ‘Umbrella’， ‘N/A’， ‘N/A’， ‘Handbag’， ‘Tie’， ‘Suitcase’， ‘Frisbee’， ‘Skis’， ‘Snowboard’， ‘Sports-Ball’， ‘Kite’， ‘Baseball Bat’， ‘Baseball Glove’， ‘Skateboard’， ‘Surfboard’， ‘Tennis Racket’， ‘Bottle’， ‘N/A’， ‘Wine Glass’， ‘Cup’， ‘Fork’， ‘Knife’， ‘Spoon’， ‘Bowl’， ‘Banana’， ‘Apple’， ‘Sandwich’， ‘Orange’， ‘Broccoli’， ‘Carrot’， ‘Hot-Dog’， ‘Pizza’， ‘Donut’， ‘Cake’， ‘Chair’， ‘Couch’， ‘Potted Plant’， ‘Bed’， ‘N/A’， ‘Dining Table’， ‘N/A’，‘N/A’， ‘Toilet’， ‘N/A’， ‘TV’， ‘Laptop’， ‘Mouse’， ‘Remote’， ‘Keyboard’， ‘Cell-Phone’， ‘Microwave’， ‘Oven’， ‘Toaster’， ‘Sink’， ‘Refrigerator’， ‘N/A’， ‘Book’， ‘Clock’， ‘Vase’， ‘Scissors’， ‘Teddy-Bear’， ‘Hair-Dryer’， ‘Toothbrush’］

如果我們想輸出不同顏色的邊框，我們可以手動(dòng)定義我們想要的RGB格式的顏色

COLORS = ［ ［0.000， 0.447， 0.741］， ［0.850， 0.325， 0.098］， ［0.929， 0.694， 0.125］， ［0.494， 0.184， 0.556］， ［0.466， 0.674， 0.188］， ［0.301， 0.745， 0.933］ ］

格式化輸出

我們還需要重新格式化模型的輸出。給定一個(gè)轉(zhuǎn)換后的圖像，模型將輸出一個(gè)字典，包含100個(gè)預(yù)測(cè)類的概率和100個(gè)預(yù)測(cè)邊框。每個(gè)包圍框的形式為（x， y， w， h），其中（x，y）為包圍框的中心（包圍框是單位正方形［0，1］×［0，1］）， w， h為包圍框的寬度和高度。因此，我們需要將邊界框輸出轉(zhuǎn)換為初始和最終坐標(biāo)，并重新縮放框以適應(yīng)圖像的實(shí)際大小。下面的函數(shù)返回邊界框端點(diǎn)：

# Get coordinates （x0， y0， x1， y0） from model output （x， y， w， h）def get_box_coords（boxes）： x， y， w， h = boxes.unbind（1） x0， y0 = （x - 0.5 * w）， （y - 0.5 * h） x1， y1 = （x + 0.5 * w）， （y + 0.5 * h） box = ［x0， y0， x1， y1］ return torch.stack（box， dim=1）

我們還需要縮放了框的大小。下面的函數(shù)為我們做了這些：

# Scale box from ［0，1］x［0，1］ to ［0， width］x［0， height］def scale_boxes（output_box， width， height）： box_coords = get_box_coords（output_box） scale_tensor = torch.Tensor（ ［width， height， width， height］）.to（ torch.cuda.current_device（）） return box_coords * scale_tensor

現(xiàn)在我們需要一個(gè)函數(shù)來(lái)封裝我們的目標(biāo)檢測(cè)pipeline。下面的detect函數(shù)為我們完成了這項(xiàng)工作。

# Object Detection Pipelinedef detect（im， model， transform）： device = torch.cuda.current_device（） width = im.size［0］ height = im.size［1］ # mean-std normalize the input image （batch-size： 1） img = transform（im）.unsqueeze（0） img = img.to（device） # demo model only support by default images with aspect ratio between 0.5 and 2 assert img.shape［-2］ 《= 1600 and img.shape［-1］ 《= 1600， # propagate through the model outputs = model（img） # keep only predictions with 0.7+ confidence probas = outputs［‘pred_logits’］.softmax（-1）［0， ：， ：-1］ keep = probas.max（-1）.values 》 0.85 # convert boxes from ［0; 1］ to image scales bboxes_scaled = scale_boxes（outputs［‘pred_boxes’］［0， keep］， width， height） return probas［keep］， bboxes_scaled

現(xiàn)在，我們需要做的是運(yùn)行以下程序來(lái)獲得我們想要的輸出：

probs， bboxes = detect（image， detr， transform）

繪制結(jié)果

現(xiàn)在我們有了檢測(cè)到的目標(biāo)，我們可以使用一個(gè)簡(jiǎn)單的函數(shù)來(lái)可視化它們。

# Plot Predicted Bounding Boxesdef plot_results（pil_img， prob， boxes，labels=True）： plt.figure（figsize=（16，10）） plt.imshow（pil_img） ax = plt.gca（） for prob， （x0， y0， x1， y1）， color in zip（prob， boxes.tolist（）， COLORS * 100）： ax.add_patch（plt.Rectangle（（x0， y0）， x1 - x0， y1 - y0， fill=False， color=color， linewidth=2）） cl = prob.argmax（） text = f‘{CLASSES［cl］}： {prob［cl］：0.2f}’ if labels： ax.text（x0， y0， text， fontsize=15， bbox=dict（facecolor=color， alpha=0.75）） plt.axis（‘off’） plt.show（）

現(xiàn)在可以可視化結(jié)果：

plot_results（image， probs， bboxes， labels=True）

英文原文：https://medium.com/swlh/object-detection-with-transformers-437217a3d62e

編輯：jq