前言
最近在看DETR的源码，断断续续看了一星期左右，把主要的模型代码理清了。一直在考虑以什么样的形式写一写DETR的源码解析。考虑的一种形式是像之前写的YOLOv5那样的按文件逐行写，一种是想把源码按功能模块串起来。考虑了很久还是决定按第二种方式，一是因为这种方式可能会更省时间，另外就是也方便我整体再理解一下吧。
我觉得看代码就是要看到能把整个模型分功能拆开，最后再把所有模块串起来，这样才能达到事半功倍。
另外一点我觉得很重要的是：拿到一个开源项目代码，要有马上配置环境能够正常运行Debug，并且通过解析train.py马上找到主要模型相关的内容，然后着重关注模型方面的解析，像一些日志、计算mAP、画图等等代码，完全可以不看，可以省很多时间，所以以后我讲解源码都会把无关的代码完全剥离，不再讲解，全部精力关注模型、改进、损失等内容。
这一节主要讲一下DETR的Backbone部分，包括CNN和位置编码两个模块的代码。主要涉及models/backbone.py和models/position_encoding.py两个文件。
Github注释版源码：HuKai97/detr-annotations
一、Backbone整体结构
整个Backbone主要包括CNN特征提取和位置编码两个部分。代码还是比较简单的，下面开始解析源码。
首先是调用models/Backbone.py中的build_backbone函数创建Backbone:
def build_backbone(args):
# 搭建backbone
# 位置编码 PositionEmbeddingSine()
position_embedding = build_position_encoding(args)
train_backbone = args.lr_backbone > 0 # 是否需要训练backbone True
return_interm_layers = args.masks # 是否需要返回中间层结果 目标检测False 分割True
# 生成backbone resnet50
backbone = Backbone(args.backbone, train_backbone, return_interm_layers, args.dilation)
# 将backbone输出与位置编码相加 0: backbone 1: PositionEmbeddingSine()
model = Joiner(backbone, position_embedding)
model.num_channels = backbone.num_channels # 512
return model
这里首先调用build_position_encoding函数生成正余弦位置编码position_embedding：[bs,256,H/32, W/32]，其中256前128是y方向位置编码，后128是x方向位置编码；再调用Backbone类生成ResNet50对输入数据进行特征提取得到特征图[bs,2048,H/32, W/32]。最后Joiner将两者合并存储起来，方便后续使用。
一、CNN-Backbone
创建ResNet50，先调用Backbone类：
class Backbone(BackboneBase):
"""ResNet backbone with frozen BatchNorm."""
def __init__(self, name: str,
train_backbone: bool,
return_interm_layers: bool,
dilation: bool):
# 直接掉包 调用torchvision.models中的backbone
backbone = getattr(torchvision.models, name)(
replace_stride_with_dilation=[False, False, dilation],
pretrained=is_main_process(), norm_layer=FrozenBatchNorm2d)
# resnet50 2048
num_channels = 512 if name in ('resnet18', 'resnet34') else 2048
super().__init__(backbone, train_backbone, num_channels, return_interm_layers)
这个类是继承自BackboneBase类的，而且CNN直接调用的就是torchvision.models中的模型，所以直接看BackboneBase类：
class BackboneBase(nn.Module):
def __init__(self, backbone: nn.Module, train_backbone: bool, num_channels: int, return_interm_layers: bool):
super().__init__()
for name, parameter in backbone.named_parameters():
# layer0 layer1不需要训练 因为前面层提取的信息其实很有限 都是差不多的 不需要训练
if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:
parameter.requires_grad_(False)
# False 检测任务不需要返回中间层
if return_interm_layers:
return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"}
else:
return_layers = {'layer4': "0"}
# 检测任务直接返回layer4即可 执行torchvision.models._utils.IntermediateLayerGetter这个函数可以直接返回对应层的输出结果
self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)
self.num_channels = num_channels
def forward(self, tensor_list: NestedTensor):
"""
tensor_list: pad预处理之后的图像信息
tensor_list.tensors: [bs, 3, 608, 810]预处理后的图片数据 对于小图片而言多余部分用0填充
tensor_list.mask: [bs, 608, 810] 用于记录矩阵中哪些地方是填充的（原图部分值为False，填充部分值为True）
"""
# 取出预处理后的图片数据 [bs, 3, 608, 810] 输入模型中 输出layer4的输出结果 dict '0'=[bs, 2048, 19, 26]
xs = self.body(tensor_list.tensors)
# 保存输出数据
out: Dict[str, NestedTensor] = {}
for name, x in xs.items():
m = tensor_list.mask # 取出图片的mask [bs, 608, 810] 知道图片哪些区域是有效的 哪些位置是pad之后的无效的
assert m is not None
# 通过插值函数知道卷积后的特征的mask 知道卷积后的特征哪些是有效的 哪些是无效的
# 因为之前图片输入网络是整个图片都卷积计算的 生成的新特征其中有很多区域都是无效的
mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0]
# out['0'] = NestedTensor: tensors[bs, 2048, 19, 26] + mask[bs, 19, 26]
out[name] = NestedTensor(x, mask)
# out['0'] = NestedTensor: tensors[bs, 2048, 19, 26] + mask[bs, 19, 26]
return out
这个类还是在调用torchvision.models中的模型，然后再把预处理后的图片数据[bs, 3, 608, 810]和mask数据[bs, 608, 810]输入到模型中（这个图片数据是经过pad填充的数据，而mask数据就是记录这些图片哪些像素位置是pad的，为True，没用pad的真实有效数据就为False）。经过前向传播，再调用IntermediateLayerGetter函数把对应层特征图提取出来，得到原图32倍下采样的特征图[bs, 2048, 19, 26]，以及这张特征图对应的mask[bs, 19, 26]。
二、Positional Encoding
Positional Encoding 就是位置编码。这里主要是调用models/position_encoding.py中的build_position_encoding函数创建位置编码：
def build_position_encoding(args):
"""
创建位置编码
args: 一系列参数 args.hidden_dim: transformer中隐藏层的维度 args.position_embedding: 位置编码类型 正余弦sine or 可学习learned
"""
# N_steps = 128 = 256 // 2 backbone输出[bs,256,25,34] 256维度的特征
# 而传统的位置编码应该也是256维度的, 但是detr用的是一个x方向和y方向的位置编码concat的位置编码方式 这里和ViT有所不同
# 二维位置编码 前128维代表x方向位置编码 后128维代表y方向位置编码
N_steps = args.hidden_dim // 2
if args.position_embedding in ('v2', 'sine'):
# TODO find a better way of exposing other arguments
# [bs,256,19,26] dim=1时 前128个是y方向位置编码 后128个是x方向位置编码
position_embedding = PositionEmbeddingSine(N_steps, normalize=True)
elif args.position_embedding in ('v3', 'learned'):
position_embedding = PositionEmbeddingLearned(N_steps)
else:
raise ValueError(f"not supported {args.position_embedding}")
return position_embedding
可以看到，源码是实现了两种位置编码，一种是正余弦绝对位置编码，不需要额外的参数学习，另一种是可学习绝对位置编码。原论文用的是正余弦绝对位置编码，而且代码也是默认使用这个的，所以这里主要介绍PositionEmbeddingSine类：
class PositionEmbeddingSine(nn.Module):
"""
Absolute pos embedding, Sine. 没用可学习参数 不可学习 定义好了就固定了
This is a more standard version of the position embedding, very similar to the one
used 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_feats # 128维度 x/y = d_model/2
self.temperature = temperature # 常数 正余弦位置编码公式里面的10000
self.normalize = normalize # 是否对向量进行max规范化 True
if scale is not None and normalize is False:
raise ValueError("normalize should be True if scale is passed")
if scale is None:
# 这里之所以规范化到2*pi 因为位置编码函数的周期是[2pi, 20000pi]
scale = 2 * math.pi # 规范化参数 2*pi
self.scale = scale
def forward(self, tensor_list: NestedTensor):
x = tensor_list.tensors # [bs, 2048, 19, 26] 预处理后的 经过backbone 32倍下采样之后的数据 对于小图片而言多余部分用0填充
mask = tensor_list.mask # [bs, 19, 26] 用于记录矩阵中哪些地方是填充的（原图部分值为False，填充部分值为True）
assert mask is not None
not_mask = ~mask # True的位置才是真实有效的位置
# 考虑到图像本身是2维的 所以这里使用的是2维的正余弦位置编码
# 这样各行/列都映射到不同的值 当然有效位置是正常值 无效位置会有重复值 但是后续计算注意力权重会忽略这部分的
# 而且最后一个数字就是有效位置的总和，方便max规范化
# 计算此时y方向上的坐标 [bs, 19, 26]
y_embed = not_mask.cumsum(1, dtype=torch.float32)
# 计算此时x方向的坐标 [bs, 19, 26]
x_embed = not_mask.cumsum(2, dtype=torch.float32)
# 最大值规范化 除以最大值 再乘以2*pi 最终把坐标规范化到0-2pi之间
if self.normalize:
eps = 1e-6
y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) # 0 1 2 .. 127
# 2i/2i+1: 2 * (dim_t // 2) self.temperature=10000 self.num_pos_feats = d/2
dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) # 分母
pos_x = x_embed[:, :, :, None] / dim_t # 正余弦括号里面的公式
pos_y = y_embed[:, :, :, None] / dim_t # 正余弦括号里面的公式
# x方向位置编码: [bs,19,26,64][bs,19,26,64] -> [bs,19,26,64,2] -> [bs,19,26,128]
pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
# y方向位置编码: [bs,19,26,64][bs,19,26,64] -> [bs,19,26,64,2] -> [bs,19,26,128]
pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
# concat: [bs,19,26,128][bs,19,26,128] -> [bs,19,26,256] -> [bs,256,19,26]
pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
# [bs,256,19,26] dim=1时 前128个是y方向位置编码 后128个是x方向位置编码
return pos
对照公式：
我的几个关键点的理解：
- 这里是通过mask来构建位置编码的，mask中记录了特征图中每个像素位置是否是pad的，只有在为False的位置，才是有效的位置，才需要构建位置编码；
- 关于最大值规范化 ：因为正余弦编码方式，思想就是将各个位置的通过公式映射到 0~2Π 这个范围内（也可以是4Π，6Π，8Π…，因为它是一个周期函数，不过我们一般默认为2Π），所以这里在带入公式之前需要对x_embed、y_embed先进行规范化；
- 关于位置编码方式：这里之所以是把x和y分别进行位置编码（二维位置编码），而不是像transformer那样的一维位置编码。主要考虑的是transformer是应用在语言模型中的，天然就是一维的，所以一维可能更适合，而DETR是应用在图像任务中的一个目标检测框架，在图像任务中，当然二维位置编码效果可能会更好点；
- 这样，对于每个位置（x,y），其所在列对应的编码值就在通道维度的前128维，其所在行的编码值就在通道这个维度的后128维。这样这个特征图上各个位置就都对应到不同的维度的编码值了。
当然作为学习，也可以看看第二种绝对位置编码方式：可学习位置编码：
class PositionEmbeddingLearned(nn.Module):
"""
Absolute pos embedding, learned.
可以发现整个类其实就是初始化了相应shape的位置编码参数，让后通过可学习的方式学习这些位置编码参数
"""
def __init__(self, num_pos_feats=256):
super().__init__()
# nn.Embedding 相当于 nn.Parameter 其实就是初始化函数
self.row_embed = nn.Embedding(50, num_pos_feats)
self.col_embed = nn.Embedding(50, num_pos_feats)
self.reset_parameters()
def reset_parameters(self):
nn.init.uniform_(self.row_embed.weight)
nn.init.uniform_(self.col_embed.weight)
def forward(self, tensor_list: NestedTensor):
x = tensor_list.tensors
h, w = x.shape[-2:] # 特征图h w
i = torch.arange(w, device=x.device)
j = torch.arange(h, device=x.device)
x_emb = self.col_embed(i) # 初始化x方向位置编码
y_emb = self.row_embed(j) # 初始化y方向位置编码
# concat x y 方向位置编码
pos = torch.cat([
x_emb.unsqueeze(0).repeat(h, 1, 1),
y_emb.unsqueeze(1).repeat(1, w, 1),
], dim=-1).permute(2, 0, 1).unsqueeze(0).repeat(x.shape[0], 1, 1, 1)
return pos
可以发现整个类其实就是初始化了相应shape的位置编码参数，然后通过可学习的方式自己学习这些位置编码参数，代码比较简答。
Reference
官方源码: https://github.com/facebookresearch/detr
b站源码讲解: 铁打的流水线工人
知乎【布尔佛洛哥哥】: DETR 源码解读
CSDN【在努力的松鼠】源码讲解: DETR源码笔记（一）
CSDN【在努力的松鼠】源码讲解: DETR源码笔记（二）
CSDN: Transformer中的position encoding(位置编码一)
知乎CV不会灰飞烟灭-【源码解析目标检测的跨界之星DETR（一）、概述与模型推断】
知乎CV不会灰飞烟灭-【源码解析目标检测的跨界之星DETR（二）、模型训练过程与数据处理】
知乎CV不会灰飞烟灭-【源码解析目标检测的跨界之星DETR（三）、Backbone与位置编码】
知乎CV不会灰飞烟灭-【源码解析目标检测的跨界之星DETR（四）、Detection with Transformer】
知乎CV不会灰飞烟灭-【源码解析目标检测的跨界之星DETR（五）、loss函数与匈牙利匹配算法】
知乎CV不会灰飞烟灭-【源码解析目标检测的跨界之星DETR（六）、模型输出与预测生成】
441
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
