本文来源公众号“机器学习算法那些事”，仅用于学术分享，侵权删，干货满满。
原文链接：有位大佬逐模块解析了detr结构
Transformer在计算机视觉领域大方异彩，Detection Transformer（DETR）是Transformer在目标检测领域的成功应用。利用Transformer中attention机制能够有效建模图像中的长程关系（long range dependency），简化目标检测的pipeline，构建端到端的目标检测器。
objection detection可以理解为一个集合预测任务（预测一个边界框和分类标签的集合），现有的目标检测算法的流程需要在大量proposals/anchors上定义回归和分类任务，DETR则通过预测集合实现目标检测。
优点：
-
不需要预定义的先验anchor
-
不需要NMS的后处理策略
-
增加transformer的编码结构
-
通过前馈神经网络直接预测框的位置和类别
缺点：
-
DETR在大目标检测上性能是最好的，而小目标上稍差
-
基于match的loss导致学习很难收敛，难以学到最优的情况
本文结合论文和代码，逐一分析DETR的模块，DETR的模块主要有：
-
backbone模块
-
位置编码模块
-
transformer编码模块
-
transformer解码模块
-
前馈神经网络模块（FNNs）
-
匈牙利算法匹配模块
-
损失函数模块
论文的模块图：
1. bakbone模块
2. 位置编码模块
位置编码模块对对象每个像素点用对应的嵌入向量表示，嵌入向量的维度是d_model，图像每个像素点的嵌入向量包括y方向和x方向的嵌入向量，每个方向的维度是d_model//2，最后拼接y方向和x方向的维度，得到嵌入向量维度d_model。
为了更方便理解，假设输入图像经过backbone后的mask维度：(2,24,32)
由上节可知，backbone模块和位置编码的模块维度相同，因此对两者的输出相加得到transformer编码模块的Q和K向量。
3. transformer编码模块
如上节介绍transformer编码模块的Q和K向量来自backbone模块和图像像素的位置编码模块之和。如下图的绿色框表示：
V向量来自backbone模块的输出。
然后经过多头注意力机制得到输出向量
最后经过残差模块和层归一化模块得到编码输出的向量。
维度表示：
向量Q：tensor(768,2,256)
向量K：tensor(768,2,256)
向量V：tensor(768,2,256)
经过多头注意力机制得到的向量维度：tensor(768,2,256)
4. Transformer解码模块
如上图object quries本质上是N个可学习的嵌入维度参数，训练开始时初始化为0，且对应的位置编码模块随机初始化，object queries是预定义的目标查询的个数，代码中默认为100。它的意义是：根据Encoder编码的特征，Decoder将100个查询转化成100个目标，即最终预测这100个目标的类别和bbox位置。最终预测得到的shape应该为[N, 100, C]，N为Batch Num，100个目标，C为预测的100个目标的类别数+1（背景类）以及bbox位置（4个值）。
解码模块共有两个自注意机制模块。
第一个自注意机制的Q向量和K向量是object quires向量与对应位置编码模块query_pos向量相加，V向量则是object quires向量。流程如下图：
第二个自注意机制模块的Q向量来自上一个自注意机制模块的输出向量tgt相应位置编码模块query_pos向量相加，K向量来自backbone模块的输出memory向量与相应位置编码模块pos_embed向量相加，V向量来自backbone模块的输出memory向量。
5. 前馈神经网络模块（FFNs）
前馈神经网络模块对transformer结构输出后的数据完成分类+bbox预测任务，两个FFN共享transformer的输出。
-
分类的class FNN使用简单的一层线性层，接受（100,C+1）维度的向量，C+1为数据集类别数+背景的个数
-
bbox预测的FFN使用三层MLP，接受（100,256）维度输入，维度变化为（100,256）->（100,512）->（100,4）。前两层计算后都接ReLU激活函数，最后一层接sigmoid函数。最终输出维度（100,4），代表目标的4个bbox坐标（x,y,w,h）
如下图：
Decoder输出向量：tensor(6,bs,100,256)
Class输出向量：tensor(6,bs,100,C+1)，其中C为类别数
Bounding Box向量：tensor(6,Bs,100,4)
6. 匈牙利算法匹配模块
前馈神经网络模块得到预测类和预测框后，通过匈牙利算法进行二分匹配，假如有K个真实目标，那么100个预测框中就会有K个能够匹配到这K个真实目标，其他都会和“no object”匹配成功，使其在理论上每个object query都会有唯一匹配的目标，不会存在重叠，所有DETR不需要nms后处理。
匈牙利算法匹配的输入向量是损失矩阵，如预测框个数是N，真实框个数是M，那么损失矩阵的形状是（N，M），表示每个预测框与所有真实框的损失值。如下图：
匈牙利算法依据匹配的总损失值最小的原理，得到预测框和真实框的匹配项。
7. 损失函数模块
损失函数模块计算类损失和预测框损失。
-
类损失模块计算所有预测框。
-
预测框损失计算匹配预测框的L1损失和GIoU损失
DETR论文链接：https://arxiv.org/abs/2005.12872
官方代码：https://github.com/facebookresearch/detr
由于官方代码没有推理模块，本文附上推理代码inference.py
内容到这就结束了，感兴趣的可以扫码关注一下，谢谢：
以下是正确缩进后的代码：
```python
import argparse
import random
import time
from pathlib import Path
import numpy as np
import torch
from models import build_model
from PIL import Image
import os
import torchvision
from torchvision.ops.boxes import batched_nms
import cv2
def get_args_parser():
parser = argparse.ArgumentParser('Set transformer detector', add_help=False)
parser.add_argument('--lr', default=1e-4, type=float)
parser.add_argument('--lr_backbone', default=1e-5, type=float)
parser.add_argument('--batch_size', default=2, type=int)
parser.add_argument('--weight_decay', default=1e-4, type=float)
parser.add_argument('--epochs', default=300, type=int)
parser.add_argument('--lr_drop', default=200, type=int)
parser.add_argument('--clip_max_norm', default=0.1, type=float, help='gradient clipping max norm')
# Model parameters
parser.add_argument('--frozen_weights', type=str, default=None, help="Path to the pretrained model. If set, only the mask head will be trained")
# * Backbone
parser.add_argument('--backbone', default='resnet50', type=str, help="Name of the convolutional backbone to use")
parser.add_argument('--dilation', action='store_true', help="If true, we replace stride with dilation in the last convolutional block (DC5)")
parser.add_argument('--position_embedding', default='sine', type=str, choices=('sine', 'learned'), help="Type of positional embedding to use on top of the image features")
# * Transformer
parser.add_argument('--enc_layers', default=6, type=int, help="Number of encoding layers in the transformer")
parser.add_argument('--dec_layers', default=6, type=int, help="Number of decoding layers in the transformer")
parser.add_argument('--dim_feedforward', default=2048, type=int, help="Intermediate size of the feedforward layers in the transformer blocks")
parser.add_argument('--hidden_dim', default=256, type=int, help="Size of the embeddings (dimension of the transformer)")
parser.add_argument('--dropout', default=0.1, type=float, help="Dropout applied in the transformer")
parser.add_argument('--nheads', default=8, type=int, help="Number of attention heads inside the transformer's attentions")
parser.add_argument('--num_queries', default=100, type=int, help="Number of query slots")
parser.add_argument('--pre_norm', action='store_true')
# * Segmentation
parser.add_argument('--masks', action='store_true', help="Train segmentation head if the flag is provided")
# Loss
parser.add_argument('--no_aux_loss', dest='aux_loss', action='store_false', help="Disables auxiliary decoding losses (loss at each layer)")
# * Matcher
parser.add_argument('--set_cost_class', default=1, type=float, help="Class coefficient in the matching cost")
parser.add_argument('--set_cost_bbox', default=5, type=float, help="L1 box coefficient in the matching cost")
parser.add_argument('--set_cost_giou', default=2, type=float, help="giou box coefficient in the matching cost")
# * Loss coefficients
parser.add_argument('--mask_loss_coef', default=1, type=float)
parser.add_argument('--dice_loss_coef', default=1, type=float)
parser.add_argument('--bbox_loss_coef', default=5, type=float)
parser.add_argument('--giou_loss_coef', default=2, type=float)
parser.add_argument('--eos_coef', default=0.1, type=float, help="Relative classification weight of the no-object class")
# dataset parameters
parser.add_argument('--dataset_file', default='coco')
parser.add_argument('--coco_path', type=str, default="coco")
parser.add_argument('--coco_panoptic_path', type=str)
parser.add_argument('--remove_difficult', action='store_true')
# 检测的图像路径
parser.add_argument('--source_dir', default='demo/images', help='path where to save, empty for no saving')
# 检测结果保存路径
parser.add_argument('--output_dir', default='demo/outputs', help='path where to save, empty for no saving')
parser.add_argument('--device', default='cuda:0', help='device to use for training / testing')
parser.add_argument('--seed', default=42, type=int)
# resnet50对应的权重文件
parser.add_argument('--resume', default='demo/weights/detr-r50-e632da11.pth', help='resume from checkpoint')
parser.add_argument('--start_epoch', default=0, type=int, metavar='N', help='start epoch')
parser.add_argument('--eval', default="True")
parser.add_argument('--num_workers', default=2, type=int)
# distributed training parameters
parser.add_argument('--world_size', default=1, type=int, help='number of distributed processes')
parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training')
return parser
def box_cxcywh_to_xyxy(x):
x_c, y_c, w, h = x.unbind(1)
b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)]
return torch.stack(b, dim=1)
def rescale_bboxes(out_bbox, size):
img_w, img_h = size
b = box_cxcywh_to_xyxy(out_bbox) # out_bbox(100,4),cx,cy,w,h,b(100,4),x1,y1,x2,y2
b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)
return b
def filter_boxes(scores, boxes, confidence=0.7, apply_nms=True, iou=0.5):
keep = scores.max(-1).values > confidence # tensor(100)
scores, boxes = scores[keep], boxes[keep] # (27,91),(27,4)
if apply_nms:
top_scores, labels = scores.max(-1)
keep = batched_nms(boxes, top_scores, labels, iou)
scores, boxes = scores[keep], boxes[keep]
return scores, boxes
# COCO classes
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 drier', 'toothbrush'
]
def plot_one_box(x, img, color=None, label=None, line_thickness=1):
tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness
color = color or [random.randint(0, 255) for _ in range(3)]
c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
if label:
tf = max(tl - 1, 1) # font thickness
t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled
cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
def main(args):
print(args)
device = torch.device(args.device)
model, criterion, postprocessors = build_model(args)
checkpoint = torch.load(args.resume, map_location='cpu')
model.load_state_dict(checkpoint['model'], False)
model.to(device)
n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
print("parameters:", n_parameters)
image_Totensor = torchvision.transforms.ToTensor()
image_file_path = os.listdir(args.source_dir)
for image_item in image_file_path:
print("inference_image:", image_item)
image_path = os.path.join(args.source_dir, image_item)
image = Image.open(image_path)
image_tensor = image_Totensor(image) # tensor(c,h,w)
image_tensor = torch.reshape(image_tensor, [-1, image_tensor.shape[0], image_tensor.shape[1], image_tensor.shape[2]])
image_tensor = image_tensor.to(device)
time1 = time.time()
inference_result = model(image_tensor) # 'pred_logits':(1,100,92),'pred_boxes':(1,100,4),list(5)
time2 = time.time()
print("inference_time:", time2 - time1)
probas = inference_result['pred_logits'].softmax(-1)[0, :, :-1].cpu() # 得到除了背景类的前景，(100,91)
bboxes_scaled = rescale_bboxes(inference_result['pred_boxes'][0,].cpu(), (image_tensor.shape[3], image_tensor.shape[2])) # (100,4)
scores, boxes = filter_boxes(probas, bboxes_scaled) # (100,91),(100,4)
scores = scores.data.numpy()
boxes = boxes.data.numpy()
for i in range(boxes.shape[0]):
class_id = scores[i].argmax()
label = CLASSES[class_id]
confidence = scores[i].max()
text = f"{label} {confidence:.3f}"
print(text)
image = np.array(image)
plot_one_box(boxes[i], image, label=text)
cv2.imshow("images", cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
cv2.waitKey()
image = Image.fromarray(image)
image.save(os.path.join(args.output_dir, image_item))
if __name__ == '__main__':
parser = argparse.ArgumentParser('DETR training and evaluation script', parents=[get_args_parser()])
args = parser.parse_args()
if args.output_dir:
Path(args.output_dir).mkdir(parents=True, exist_ok=True)
main(args)THE END !
文章结束，感谢阅读。您的点赞，收藏，评论是我继续更新的动力。大家有推荐的公众号可以评论区留言，共同学习，一起进步。
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
