目标检测DETR理解: 从网络前向张量的变化开始

文章 https://arxiv.org/pdf/2005.12872.pdf
官方git仓库 https://github.com/facebookresearch/detr

1. 基础原理

参考 https://zhuanlan.zhihu.com/p/308301901

2. 输入batch图像到输出的张量变化

先看纵览图:

分3部分解析,
第一部分: transformer前的预操作. 包括backbone提取特征、位置编码pos_embed(不可训练,比如sincos编码)生成、图像mask生成、查询编码query_embed(重要, 可训练, 表示全局的图像中各个目标之间的关系)生成。
注意到query_embed的shape是[num_query, c], 这个c就是256, 和feat以及pos_embed一样; 而num_query表示1张图中最多的目标个数, 类似nlp中1个句子的最大长度.


第二部分: 使用encoder对图像特征进一步提取各个区域间的关系信息. 一个encoder(编码器)包括6个encoder layer(编码层)


第三部分: 使用decoder, 在encoder提取的信息 + query_embed(重要,第一步就说了,是可学习的) + pos_embed(位置编码)基础上, 进行特征的解码
注意, encoder中输入的qkv这3个的shape都一样, 都是[h*w, bs, 256], 并且在整个encoder(6个编码层)过程中都一样, 其中第一个坐标轴表示length, 相当于nlp中的翻译任务的句子长度, 
而decoder中, 每个decoder layer中, 第一个self_attn(注意力层)的qkv都一样, 都是q[num_query, bs, 256], 而第二个cross_attn layer(一样,也是注意力层)就不一样了.
q的length是num_query, kv的length是h*w. 虽然length不一样, 但是这个层的输出shape就是[num_query(from q), bs, 256].详细过程可以看附录.


得到decoder的输出dec_output[h*w, bs, 256]之后, 后续再输给两个分类和回归分支(这个就和一般的目标检测一样), 来得到最后100个目标框

3. 正负样本构建与loss

暂时pass 参考 https://zhuanlan.zhihu.com/p/308301901

4. 测试阶段解析出bboxes

暂时pass 参考 https://zhuanlan.zhihu.com/p/308301901
 

5.附录

5.1 attention层计算过程

为了方便这里只考虑qkv, 不考虑pos_embed位置编码.
这里输入格式是[bs, length, dim], 而pytorch的torch.nn.MultiheadAttention层默认参数batch_first=False, 所以默认输入是[length, bs, dim], 与这里有区别
运行下面demo就知道对于attention层的输入qkv: 
# q的length可以不一样, k和v的length必须一样
# v的dim可以不一样, q和k的dim必须一样
# 最终是q的length, v的dim

import torch
import torch.nn as nn
import torch.nn.functional as F


bs = 10
N = 100 # q的序列长度
m = 512
q = torch.rand([bs, N, m])
k = torch.rand([bs, N+10, m])
v = torch.rand([bs, N+10, m])

d_model = m
n_head = 8
dim_k = 64
dim_v = 64
temperature = dim_k ** 0.5
# 如果dim_v * n_head == d_model,
# 则最后多头的输出concat起来的维度还是和原来的输入一样, 就不用另外加fc了; 否则要加


print("===")
print("这里qkv是原始输入, 都一样都是x")
print("3个坐标轴分别表示[bs, length, dim]")
print("q: ", q.shape)
print("k: ", k.shape)
print("v: ", v.shape)



## 1. 注册多头: n_head个头, q和v的维度必须是一样的
W_q = nn.Linear(d_model, n_head * dim_k, bias=False)
W_k = nn.Linear(d_model, n_head * dim_k, bias=False)
W_v = nn.Linear(d_model, n_head * dim_v, bias=False)

bs, _, _ = q.shape
q = W_q(q).view(bs, -1, n_head, dim_k)
k = W_k(k).view(bs, -1, n_head, dim_k)
v = W_v(v).view(bs, -1, n_head, dim_v)
print("===")
print("将输入与可学习矩阵相乘得到qkv")
print("q: ", q.shape)
print("k: ", k.shape)
print("v: ", v.shape)



## 2. 把batch和n_head两个维度提到前面, length和dim维度放在后面为了方面计算
q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
print("===")
print("转换一下坐标系, 准备计算注意力")
print("q: ", q.shape)
print("k: ", k.shape)
print("v: ", v.shape)



## 3. 一次性计算batch所有样本和所有head的注意力 [bs, n_head, length, length]
# [bs, n_head, length, dim_k] 矩阵乘以 [bs, n_head, dim_k, length]
# torch.matmul用法可以参考https://blog.csdn.net/qsmx666/article/details/105783610
print("===")
print("得到length维度之间相互的注意力")
print("输入q: ", q.shape)
tmpk = k.transpose(2, 3)
print("输入tmpk: ", tmpk.shape)
attn = torch.matmul(q / temperature, tmpk)
attn = F.softmax(attn, dim=-1)
print("attn: ", attn.shape)
print("这是q和k序列中两两元素对应的注意力(相关性)")



## 4. 根据注意力计算值向量: 1个步长的值向量是所有步长值向量的加权求和, 用矩阵相乘实现即可
print("===")
print("得到注意力与v的加权求和")
print("输入attn: ", attn.shape)
print("输入v: ", v.shape)
q = torch.matmul(attn, v)
print("输出q: ", q.shape)
length_q = q.shape[2]
q = q.transpose(1, 2).reshape(bs, length_q, -1)
print("===")
print("将n_head维度移到后面")
print("q: ", q.shape)

调试输出如下: