文章目录
- 端到端目标检测框架DETR
- 背景介绍
- 模型结构
- 模块解析
- 数据
- 模型结构
- 动手实现`DETR`
欢迎访问个人网络日志🌹🌹知行空间🌹🌹
端到端目标检测框架DETR
背景介绍
DETR
是Facebook AI
的Nicolas Carion
等于2020
年05
月提交的论文中提出的。
论文地址: https://arxiv.org/abs/2005.12872
开源代码: https://github.com/facebookresearch/detr
DETR(DEtection TRansformer)
将目标检测问题看成是集合预测的问题,所谓集合预测set prediction
是指一次输出一张图像中的所有待检测对象。
DETR
使用transformer
来做目标检测,直接预测检测框到检测框中心点归一化的距离。在模型训练时,Proposal Assignment
使用的算法是一对一的匈牙利算法,通过query
的方式获取最后的输出。以上介绍的策略,使得DETR
实现了目标检测算法的端到端训练,不需要使用NMS
和先验anchor
。
模型结构
从上面这个图可以看到DETR
的架构相当简单,输入一张图像,直接输出的就是所有的检测框,不需要复杂的编解码,不需要NMS
。
模块解析
数据
官方源码中数据定义在CocoDetection
类中,这个类继承自torchvision.datasets.CocoDetection
只需要传入COCO
格式数据集的图像和json
标注文件即可,
COCO
格式数据集文件夹路径:
.
├── annotations
│ ├── train.json
│ └── val.json
└── images├── train└── val
其中,标签文件bounding box
的格式为:
left top width height
在CoCoDetection
类中有一个self.prepare
属性,这是一个函数,其中会将ltwh
格式的检测框变换成x1y1x2y2
格式的检测框。
在DETR
源码中使用的变换函数不是从torchvision
中导入的,而是自定义的,可以看到在Normalize
中,不仅处理了图像数据,还将检测框从x1y1x2y2
格式变换成了cxcywh
格式,并相对于图像的宽高进行了归一化,其值变换到了[0,1]
。
class Normalize(object):def __init__(self, mean, std):self.mean = meanself.std = stddef __call__(self, image, target=None):image = F.normalize(image, mean=self.mean, std=self.std)if target is None:return image, Nonetarget = target.copy()h, w = image.shape[-2:]if "boxes" in target:boxes = target["boxes"]boxes = box_xyxy_to_cxcywh(boxes)boxes = boxes / torch.tensor([w, h, w, h], dtype=torch.float32)target["boxes"] = boxesreturn image, target
模型结构
DETR
的模型结构其实很简单,先是将图像输入到几层卷积神经网络中得到特征图feature map
,然后使用src = src.flatten(2).permute(2, 0, 1)
将特征图WH
维度拉平将图像变换成长度为L=W*H
的序列数据。
根据序列的长度和每个Token
的通道数生成位置编码。
将feature map
生成的序列和位置编码信息相加作为transformer
的输入src
。
除了输入的特征序列之外,还输入了图像数据的掩码src_mask
。原因是因为一个batch
输入的图像宽高不一定相同,源码中的处理方式是取一个batch
中尺寸最大的图像尺寸,其余图像往右下方向补0,最后变成尺寸一致的图像用于计算。这是为了避免padding-0
参与计算,需要将src_mask
输入到transformer
中。
DETR
使用的位置编码是针对图像的带mask
的二维位置编码,
class PositionEmbeddingSine(nn.Module):"""This is a more standard version of the position embedding, very similar to the oneused by the Attention is all you need paper, generalized to work on images."""def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):super().__init__()self.num_pos_feats = num_pos_featsself.temperature = temperatureself.normalize = normalizeif scale is not None and normalize is False:raise ValueError("normalize should be True if scale is passed")if scale is None:scale = 2 * math.piself.scale = scaledef forward(self, tensor_list: NestedTensor):x = tensor_list.tensorsmask = tensor_list.maskassert mask is not Nonenot_mask = ~masky_embed = not_mask.cumsum(1, dtype=torch.float32)x_embed = not_mask.cumsum(2, dtype=torch.float32)if self.normalize:eps = 1e-6y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scalex_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scaledim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)pos_x = x_embed[:, :, :, None] / dim_tpos_y = y_embed[:, :, :, None] / dim_tpos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)return pos
其在x/y
单个方向上使用位置编码的方法同标准的transformer
,然后再将x,y
上的两个位置编码分别进行了合并。
DETR
源码中使用的transformer
和torch.nn.Transformer
也不太一样。
DETR
的transformer
中将位置编码信息输入到编码器和解码器的每一层,在encoder
中将pos
加在输入的feature
上组成q
和k
,
class Encoder:...def forward_post(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
在decoder
中将pos
加在了encoder
的输出memory
作为k
的值,query_pos
和tgt
相加的值作为q
来计算多头注意力:
class Decoder:def forward_post(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
在DETR
中实现的transformer
中还将每层decoder
输出都保存下来以计算检测框,用来辅助训练。
class DETRTransformerDecoder():...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):output = tgtintermediate = []for layer in self.layers:output = layer(output, memory, tgt_mask=tgt_mask,memory_mask=memory_mask,tgt_key_padding_mask=tgt_key_padding_mask,memory_key_padding_mask=memory_key_padding_mask,pos=pos, query_pos=query_pos)if self.return_intermediate:intermediate.append(self.norm(output))if self.norm is not None:output = self.norm(output)if self.return_intermediate:intermediate.pop()intermediate.append(output)if self.return_intermediate:return torch.stack(intermediate)return output.unsqueeze(0)
将transformer
输出的特征输入到计算评分和检测框的两支多层感知积网络中就能预测检测框了:
class DETR:...def forward(self, x):hs = self.transformer(self.input_proj(src), mask, self.query_embed.weight, pos[-1])[0] # shape: [BATCH, NUM_QUERY, D_MODEL]outputs_class = self.class_embed(x)outputs_coord = self.box_embed(x).sigmoid()
以上就是模型的整体结构。
模型输出的num_query
个预测框和真值框之间的匹配通过匈牙利算法来实现。匈牙利算法会实现预测框和真值框的一对一匹配,避免了对同个对象生成重复的检测框。在使用anchor
的检测算法中,为了减轻候选框中正样本和负样本不平衡的问题,通常会使用多个proposal box
来预测一个对象,以提升算法的召回率,代价是预测推理时也会对一个对象生成多个预测框,需要使用NMS
算法进行处理。
标签匹配使用的代价包括三部分,分别是分类代价,检测框回归相关的L1
距离和GIoU
。
import torchclass HungarianMatcher(torch.nn.Module):...@torch.no_grad()def forward(self, outputs, targets):...cost_class = -out_prob[:, tgt_ids]cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1)cost_giou = -giou(cxcywh2x1y1x2y2(out_bbox), cxcywh2x1y1x2y2(tgt_bbox))all_cost = self.cost_class * cost_class + \self.cost_bbox * cost_bbox + \self.cost_giou * cost_giou
最后是模型训练时使用的损失函数,对于目标检测任务,DETRLoss
包含2
部分,分别是标签类别损失和检测框回归的L1
损失和GIoU
损失。
loss_ce = torch.nn.functional.cross_entropy(pred_logits.transpose(1, 2),target_classes_all, self.empty_weight)loss_bbox = torch.nn.functional.l1_loss(src_boxes, target_boxes, reduction='none')
losses = {}losses["loss_bbox"] = loss_bbox.sum() / num_boxes
loss_giou = 1 - torch.diag(giou(cxcywh2x1y1x2y2(src_boxes),cxcywh2x1y1x2y2(target_boxes)))
losses['loss_giou'] = loss_giou.sum() / num_boxes
动手实现DETR
DETR
的架构如此简洁,不需要太多的trick
,参考DETR
源码,很容易自己动手实现DETR
目标检测算法。具体的实现见:
https://gitee.com/lx_r/object_detection_task/tree/main/detection/detr
运行程序会自动生成训练数据开始训练,若平台有GPU
会自动调用GPU
训练,如果没有GPU
会使用CPU
训练。
上面的实现中,与原始代码有些许不同:
- 1)使用的是
torch.nn
中的transformer
,pos
没有加到encoder
的输出memory
上 - 2)
torch.nn
中的transformer
只给出了最后一层decoder
上的输出,没有给出其他层decoder
上的输出,所有没有使用辅助损失训练 - 3)输入的是相同尺寸的方形图像,没有使用输入掩码
欢迎访问个人网络日志🌹🌹知行空间🌹🌹