【CV学习笔记】onnx篇之DETR

1、摘要

本次学习内容主要学习了DETR的网络结构、损失函数等知识，明白了DETR是如何做到了端到端的检测，确实是一个十分优雅的框架，同时将DETR利用onnxtime进行推理，对于transformer的理解进一步加深了。

DETR学习链接:https://www.bilibili.com/video/BV1GB4y1X72R?spm_id_from=333.337.search-card.all.click

DETR官方地址:https://github.com/facebookresearch/detr

个人学习地址:https://github.com/Rex-LK/tensorrt_learning

2、DETR

2.1、简介

DETR是transformer在目标检测领域内的里程碑式的工作，主要实现了端到端的目标检测，避免了计算anchor和nms操作，其网络结构也十分直接明了，下图为论文中详细绘制的DETR网络结构图，大致可以分为如下四个步骤:

1、利用CNN提取特征图

2、encoder用于学习全局的特征

3、decoder生成预测框

4、训练时，采用二分图匹配的方式将ground truth框和预测框做匹配，并计算loss。预测时，直接将第三步生成的预测框中阈值低于0.7的过滤掉。

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-y61oGTQH-1655255959113)(Screenshot%20from%202022-06-12%2020-48-11.png)]

原文中图片的输入尺寸是3×800×1066，通过卷积提取特征之后得到了2048×25×34的特征图，特征层尺寸为原图的1/32，然后将2048个特征层映射到256，变成256×25×34，同时positional ecodeing 维度也为256×25×34，位置编码与特征层相加之后，然后将25*34展平，最后得到850×256的特征向量，输入到transformer中，然后通过6个encoder后得到850×256的全局特征，然后输入到decoder中。

在decoder中加入了 object queries，是一个可学习的向量，维度为100×256，其中100代表预测100个预测，然后再将每层的 object querries与 每层的850×256 特征层反复做自注意力操作，就是将object querries 当做querry，将每层decoder得到的输出作为key，最终得到了一个100*×256的特征，然后利用FFN预测出物体的类别以及xywh，利用预测的100个框和ground truth 做最优匹配，采用匈牙利算法计算最后的目标函数。

其中在decoder中第一层没有object quireies，后面五层才有，主要是为了移除冗余的框，在object quireies通信之后，就可以知道其他每个query预测出什么框，然后尽量不要去重复这个框，似的最后不需要进行nms操作。在最后计算loss的时候，为了加速收敛并训练的更稳定，在每一个decoder后(共6个)都加了auxiliary loss。

下面通过论文中给出的推理代码

import torch
from torch import nn
from torchvision.models import resnet50

class DETR(nn.Module):
	def __init__(self, num_classes, hidden_dim, nheads,num_encoder_layers, num_decoder_layers):
		super().__init__()
        #resnet50提取图片特征
        self.backbone = nn.Sequential(*list(resnet50(pretrained=True).children())[:-2])
        #将2048个特征层映射到256
        self.conv = nn.Conv2d(2048, hidden_dim, 1)
        #encoder and decoder
        self.transformer = nn.Transformer(hidden_dim, nheads,num_encoder_layers, num_decoder_layers)
        #类别预测
        self.linear_class = nn.Linear(hidden_dim, num_classes + 1)
        #框的预测
        self.linear_bbox = nn.Linear(hidden_dim, 4)
        #object_queries 100×256
        self.query_pos = 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, inputs):
        #第一步提取特征
        #3*800*1066 -> 2048×25×34
        x = self.backbone(inputs)
        
        #256×25×34
        h = self.conv(x)
        H, W = h.shape[-2:]
        # 位置编码
        pos = 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 + h.flatten(2).permute(2, 0, 1),self.query_pos.unsqueeze(1))
        # 100×256的特征
        return self.linear_class(h), self.linear_bbox(h).sigmoid()
detr = DETR(num_classes=91, hidden_dim=256, nheads=8, num_encoder_layers=6, num_decoder_layers=6)
detr.eval()
inputs = torch.randn(1, 3, 800, 1066)
logits, bboxes = detr(inputs)

2.2、facebook官方源码学习

拿到官方源码的时候其实是很懵的，不知道从哪下手，后来经过一段时间的摸索，可以在我的仓库中找到mypredict.py，来实现一个简单的推理。

 #初始化模型
 detr = detr_resnet50()

其中模型初始化调用的_make_detr这个函数,其中backbone采用的resnet50，注意这里输出的特征层的尺寸大小为2048×50×67，是原文的2倍.

def _make_detr(backbone_name: str, dilation=False, num_classes=91, mask=False):
    hidden_dim = 256
    backbone = Backbone(backbone_name, train_backbone=False, return_interm_layers=mask, dilation=dilation)
    pos_enc = PositionEmbeddingSine(hidden_dim // 2, normalize=True)
    backbone_with_pos_enc = Joiner(backbone, pos_enc)
    backbone_with_pos_enc.num_channels = backbone.num_channels
    transformer = Transformer(d_model=hidden_dim, return_intermediate_dec=True)
    detr = DETR(backbone_with_pos_enc, transformer, num_classes=num_classes, num_queries=100)
    if mask:
        return DETRsegm(detr)
    return detr

在Transformer中定义了encoder和decoder

#定义一层encoder
encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward,
                                        dropout, activation, normalize_before)
encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
#定义六层encoder
self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)

#定义一层decoder
decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward,
                                        dropout, activation, normalize_before)
decoder_norm = nn.LayerNorm(d_model)
#定义六层decoder
self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm,
                                  return_intermediate=return_intermediate_dec)

然后在DETR中定义了一些基本组件

self.num_queries = num_queries
self.transformer = transformer
hidden_dim = transformer.d_model
self.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 = backbone
self.aux_loss = aux_loss

初始化模型之后，接着就是推理过程了

    def forward(self, samples: NestedTensor):
        if isinstance(samples, (list, torch.Tensor)):
            samples = nested_tensor_from_tensor_list(samples)
        features, pos = self.backbone(samples)
        #利用resnet50提取特征 3×800×1066 -> 2048×50×67
        src, mask = features[-1].decompose()
        assert mask is not None
        #然后经过encoder和decoder得到100*256的预测值
        hs = self.transformer(self.input_proj(src), mask, self.query_embed.weight, pos[-1])[0]
        #类别以及bbox
        #取第六个decoder的结果
        outputs_class = self.class_embed(hs)[-1]
        outputs_coord = self.bbox_embed(hs).sigmoid()[-1][0]
        ...

其中self.transformer的计算代码如下

    def forward(self, src, mask, query_embed, pos_embed):
        # flatten NxCxHxW to HWxNxC
        bs, c, h, w = src.shape
        #将后两个维度展平 
        # [3350, 1, 256]
        src = src.flatten(2).permute(2, 0, 1)
        # 位置编码 [3350, 1, 256]
        pos_embed = pos_embed.flatten(2).permute(2, 0, 1)
        #decoder中的object queries
        query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1)
        mask = mask.flatten(1)
        tgt = torch.zeros_like(query_embed)
        #memory  [3350, 1, 256]
        memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)
        #hs [6, 100, 1, 256] 六个decoder预测结果,预测时，取第六个decoder的结果
        hs = self.decoder(tgt, memory, memory_key_padding_mask=mask,
                          pos=pos_embed, query_pos=query_embed)
        return hs.transpose(1, 2), memory.permute(1, 2, 0).view(bs, c, h, w)

查看其中decoder的代码，主要是在TransformerDecoderLayer这个类中

    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)
        # 对应每除了第一个decoder，其余每个decoder都与objectquerries进行注意力计算
        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的推理过程有了一个比较直观的了解，总的来说，推理过程十分简洁，无非还是分为如下四个步骤

1、cnn提取图像特征

2、encoder提取全局特征

3、decoder生成预测框

4、筛选预测框

2.3、onnxruntime

2.3.1 export_onnx

在demo/detr-mian/mypredict.py中，包含了导出onnx以及onnx-simplify的方法，为了将部分后处理代码放到onnx中，对后处理代码进行了如下改写

hs = self.transformer(self.input_proj(src), mask, self.query_embed.weight, pos[-1])[0]
outputs_class = self.class_embed(hs)[-1]
outputs_coord = self.bbox_embed(hs).sigmoid()[-1][0]
probas = outputs_class.softmax(-1)[0, :, :-1]
pred = torch.cat((probas,outputs_coord),1)
pred = pred.unsqueeze(0)
return pred

通过netron来查看导出的onnx是否存在问题，发现导出的onnx的输出维度为1×100×95，为100个框的类别以及xywh，说明没有问题

导出onnx后，可以使用onnxruntime来检测导出onnx的正确性，运行infer-onnxruntime.py，如果结果与mypredict.py显示的结果一致，那么就说明导出的onnx正确。

if __name__ == "__main__":

    data_transform = transforms.Compose([
        transforms.Resize(800),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    img_path = "demo.jpg"
    img_o = Image.open(img_path)
    img = data_transform(img_o).unsqueeze(0)
    image_input = img.numpy()
    session = onnxruntime.InferenceSession("detr_sim.onnx", providers=["CPUExecutionProvider"])
    pred = session.run(["predict"], {"image": image_input})[0]
    scores = torch.from_numpy(pred[0][:,0:91])
    bboxes = torch.from_numpy(pred[0][:,91:])
    keep = scores.max(-1).values > 0.7
    scores = scores[keep]
    bboxes = bboxes[keep]
    print(bboxes)
    fin_bboxes = rescale_bboxes(bboxes, img_o.size)
    plot_results(img_o, scores, fin_bboxes)

可以看出，detr的后处理方式还是很简单的，由于这里转engine还有些许问题，等之后解决了这个问题之后，再进行tensorrt加速，用一张图来看看预测效果把。

3、总结

本次学习了detr的网络结构，了解了端到端的预测机制，阅读了detr的源码，受益匪浅，对transformer有了进一步的了解，只是遗憾的是暂时没能进行tensorrt加速，后续希望能解决这个问题。