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
欢迎正在学习或者想学的CV的同学进群一起讨论与学习，v:Rex1586662742，q群:468713665
2、DETR
2.1、简介
DETR是transformer在目标检测领域内的里程碑式的工作，主要实现了端到端的目标检测，避免了计算anchor和nms操作，其网络结构也十分直接明了，下图为论文中详细绘制的DETR网络结构图，大致可以分为如下四个步骤:
1、利用CNN提取特征图
2、encoder用于学习全局的特征
3、decoder生成预测框
4、训练时，采用二分图匹配的方式将ground truth框和预测框做匹配，并计算loss。预测时，直接将第三步生成的预测框中阈值低于0.7的过滤掉。
原文中图片的输入尺寸是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加速，后续希望能解决这个问题。
3956
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
