前言
JDE可不是一个模型的名字，而是一类追踪算法的总称，全称叫做
Jointly learns the Detector and Embedding model (JDE)
什么意思呢？我们之前讨论的一些多目标追踪模型，比如SORT和DeepSORT，都是2015-2018年常见的MOT范式，也就是tracking by detection 。
该类范式因为通俗易懂，且表现出了不俗的追踪精度，在2015年到2018年，一度成为MOT的主流范式。该范式首先通过检测器（detector）检测出画面中物体所在的检测框，然后根据物体检测框移动的规律（运动特征）和检测框中物体的外观特征（通常通过一个ReID网络抽取一个低维的向量，叫做embedding向量）来进行前后帧同一物体的匹配，从而实现多目标追踪。
该类范式将MOT分为了两步，即
- 目标检测
- 特征提取与物体关联 该类方法检测与特征提取是分开的，所以又被称为SDE
Separate Detection and Embedding (SDE)
SDE存在的最大缺点就是速度慢，因为将物体检测和（外观）特征提取分开，检测速度自然就下去了。
那么作者考虑到这点不足，提出了JDE范式。论文和代码链接如下：
SDR、Two-stage与JDE
作者对比了前面方法：
- SDE（Separate Detection and Embedding）就是两阶段法。检测器和Reid模块是独立开来的，先检测后识别（类似Faster-RCNN很像），这些方法需要分为两个大部分操作，两部分不相互干扰，精度高，但是耗时长。
- 这里的two-stage实际上是一种端到端的方法，只是不是利用检测的最后的结果，而是利用两阶段检测法中FPN生成的框来做Embedding，这样做可以共享部分特征，减少计算量。但是检测与生成特征部分本质上仍然是两阶段，速度依然无法达到实时。
- SDE——先从原图中检测一系列包含目标的框，再把这些框送入re-ID model，分配ID和轨迹
- Two-Stage——从原图中找到一系列ROI的feature，再给这些ROI作bbox回归、分配ID和轨迹
- JDE——从原图中直接输出检测框以及embedding，再进行分配ID和轨迹，这样检测部分和生成特征部分就融合成了单阶段。（后续的分配ID和轨迹操作不涉及神经网络）
JDE的网络结构和损失函数
JDE是如何一步输出检测结果和Embedding的？
既然作者提到该方法是基于One-stage检测器学习到物体的embedding的（代码中采用的是经典的YOLO V3模型）。那么JDE范式就应该在检测器的输出（head），多输出一个分支用来学习物体的embedding的。
果不其然，作者就是采用了这种思路。原论文中给出结构图如下
这个结构图勾勒了作者大致的想法，在Prediction head中多出来了一个分支用于输出embedding。然后使用一个多任务学习（multi-task learning） 的思路设置损失函数。看的时候觉得如此简单，但是深思了下，发现问题没有这么简单。
第一个问题就是embedding要怎么学出来？
我们知道，理想情况下同一物体在不同的帧中，被同一跟踪标签锁定（即拥有同一track ID）。我们知道的信息就只有他们的标签索引（同一物体的track ID一致，不同物体的track ID不一样）。
那么网络在训练的过程中，应该需要对embedding进行转化，转化为足够强的语义信息，也就是这个embedding到底属于哪个track ID的物体，那么这种就需要借鉴物体分类的思路了（将每个track ID当作一个类别），所以作者引入了全连接层将embedding信息转化为track ID分类信息。
第二个问题就是这个全连结层的输出节点到底是多少？
因为刚才我们提到，要将每个track ID当作一个类别，但是这个track ID的数量十分庞大，甚至不可计数。这个输出节点应该如何设置呢？看了一圈代码，代码中设置了14455个输出节点，这个非常大呀，具体设置的依据是什么，可能和训练数据集的物体数量有关吧。（如果有知道的小伙伴，欢迎留言告诉我，谢谢）
那么实际上作者的结构图是进行了省略的，我们不妨手动补全。
这样的话就没什么大的疑问了。有关YOLO V3的结构，熟悉目标检测的都不会陌生，我们这里忽略FPN网络的结构定义，直接看predicition head的代码部分，确定上面的分析是合理的。
该predicition head的代码定义在model.py文件下的 YOLOLayer类中。定义如下：
class YOLOLayer(nn.Module):
def __init__(self, anchors, nC, nID, nE, img_size, yolo_layer):
super(YOLOLayer, self).__init__()
self.layer = yolo_layer
nA = len(anchors)
self.anchors = torch.FloatTensor(anchors)
self.nA = nA # number of anchors (4)
print('nA',nA)
self.nC = nC # number of classes (1)
print('nC', nC)
self.nID = nID # number of identities, 14455
print('nID', nID)
self.img_size = 0
self.emb_dim = nE # 512
print('nE', nE)
self.shift = [1, 3, 5]
self.SmoothL1Loss = nn.SmoothL1Loss() # for bounding box regression
self.SoftmaxLoss = nn.CrossEntropyLoss(ignore_index=-1) # foreground and background classification
self.CrossEntropyLoss = nn.CrossEntropyLoss()
self.IDLoss = nn.CrossEntropyLoss(ignore_index=-1) # loss of embedding
self.s_c = nn.Parameter(-4.15*torch.ones(1)) # -4.15
self.s_r = nn.Parameter(-4.85*torch.ones(1)) # -4.85
self.s_id = nn.Parameter(-2.3*torch.ones(1)) # -2.3
self.emb_scale = math.sqrt(2) * math.log(self.nID-1) if self.nID>1 else 1
作者简单定义了一些参数和损失函数类型。我已经将一些参数的设定值标注在后面，便于大家理解。作者定义了四个损失函数，但是只用到了三个，分别为
- self.SmoothL1Loss 用于检测框的回归
- self.SoftmaxLoss 用于前景和背景分类
- self.IDLoss 用于计算embedding的损失
当然了，作者在论文中提到，JDE是一个多任务的学习，所以在计算损失函数的时候，需要采用任务独立不确定性（task-independent uncertainty）的自动学习方案进行以上三个损失函数的求和，有关参数在上述代码中已经定义，分别为
- self.s_c
- self.s_r
- self.s_id
这在上图中也有体现。该部分的核心代码如下
def forward(self, p_cat, img_size, targets=None, classifier=None, test_emb=False):
p, p_emb = p_cat[:, :24, ...], p_cat[:, 24:, ...]
nB, nGh, nGw = p.shape[0], p.shape[-2], p.shape[-1]
if self.img_size != img_size:
create_grids(self, img_size, nGh, nGw)
if p.is_cuda:
self.grid_xy = self.grid_xy.cuda()
self.anchor_wh = self.anchor_wh.cuda()
p = p.view(nB, self.nA, self.nC + 5, nGh, nGw).permute(0, 1, 3, 4, 2).contiguous() # prediction
p_emb = p_emb.permute(0,2,3,1).contiguous()
p_box = p[..., :4]
p_conf = p[..., 4:6].permute(0, 4, 1, 2, 3) # Conf
# Training
if targets is not None:
if test_emb:
tconf, tbox, tids = build_targets_max(targets, self.anchor_vec.cuda(), self.nA, self.nC, nGh, nGw)
else:
tconf, tbox, tids = build_targets_thres(targets, self.anchor_vec.cuda(), self.nA, self.nC, nGh, nGw)
tconf, tbox, tids = tconf.cuda(), tbox.cuda(), tids.cuda()
mask = tconf > 0
# Compute losses
nT = sum([len(x) for x in targets]) # number of targets
nM = mask.sum().float() # number of anchors (assigned to targets)
nP = torch.ones_like(mask).sum().float()
if nM > 0:
lbox = self.SmoothL1Loss(p_box[mask], tbox[mask])
else:
FT = torch.cuda.FloatTensor if p_conf.is_cuda else torch.FloatTensor
lbox, lconf = FT([0]), FT([0])
lconf = self.SoftmaxLoss(p_conf, tconf)
lid = torch.Tensor(1).fill_(0).squeeze().cuda()
emb_mask,_ = mask.max(1)
# For convenience we use max(1) to decide the id, TODO: more reseanable strategy
tids,_ = tids.max(1)
tids = tids[emb_mask]
embedding = p_emb[emb_mask].contiguous()
embedding = self.emb_scale * F.normalize(embedding)
nI = emb_mask.sum().float()
if test_emb:
if np.prod(embedding.shape)==0 or np.prod(tids.shape) == 0:
return torch.zeros(0, self.emb_dim+1).cuda()
emb_and_gt = torch.cat([embedding, tids.float()], dim=1)
return emb_and_gt
if len(embedding) > 1:
logits = classifier(embedding).contiguous()
lid = self.IDLoss(logits, tids.squeeze())
# Sum loss components
loss = torch.exp(-self.s_r)*lbox + torch.exp(-self.s_c)*lconf + torch.exp(-self.s_id)*lid + \
(self.s_r + self.s_c + self.s_id)
loss *= 0.5
return loss, loss.item(), lbox.item(), lconf.item(), lid.item(), nT
else:
p_conf = torch.softmax(p_conf, dim=1)[:,1,...].unsqueeze(-1)
p_emb = F.normalize(p_emb.unsqueeze(1).repeat(1,self.nA,1,1,1).contiguous(), dim=-1)
#p_emb_up = F.normalize(shift_tensor_vertically(p_emb, -self.shift[self.layer]), dim=-1)
#p_emb_down = F.normalize(shift_tensor_vertically(p_emb, self.shift[self.layer]), dim=-1)
p_cls = torch.zeros(nB,self.nA,nGh,nGw,1).cuda() # Temp
p = torch.cat([p_box, p_conf, p_cls, p_emb], dim=-1)
#p = torch.cat([p_box, p_conf, p_cls, p_emb, p_emb_up, p_emb_down], dim=-1)
p[..., :4] = decode_delta_map(p[..., :4], self.anchor_vec.to(p))
p[..., :4] *= self.stride
return p.view(nB, -1, p.shape[-1])
稍微有点冗长，我们来一一分析。
（1）前向预测信息的划分
首先作者将输出特征图划分为三部分，分别为
- 包含embedding信息的p_emb
- 包含检测框位置信息的p_box
- 包含前景背景分类置信度的p_conf
p_emb = p_emb.permute(0,2,3,1).contiguous()
p_box = p[..., :4]
p_conf = p[..., 4:6].permute(0, 4, 1, 2, 3) # Conf
这对应上图中的
（2）监督信息的划分
其次，作者从真实的标签信息中抽取各自的监督信息，分别为
# Training
if targets is not None:
if test_emb:
tconf, tbox, tids = build_targets_max(targets, self.anchor_vec.cuda(), self.nA, self.nC, nGh, nGw)
else:
tconf, tbox, tids = build_targets_thres(targets, self.anchor_vec.cuda(), self.nA, self.nC, nGh, nGw)
tconf, tbox, tids = tconf.cuda(), tbox.cuda(), tids.cuda()
mask = tconf > 0
因为我们知道，想要高效地进行损失函数求解，最好将监督信息对应到网络输出格式。
（3）检测框回归损失和前景背景分类损失
接着作者计算检测框回归损失和前景背景分类损失。
# Compute losses
nT = sum([len(x) for x in targets]) # number of targets
nM = mask.sum().float() # number of anchors (assigned to targets)
nP = torch.ones_like(mask).sum().float()
if nM > 0:
lbox = self.SmoothL1Loss(p_box[mask], tbox[mask])
else:
FT = torch.cuda.FloatTensor if p_conf.is_cuda else torch.FloatTensor
lbox, lconf = FT([0]), FT([0])
lconf = self.SoftmaxLoss(p_conf, tconf)
lid = torch.Tensor(1).fill_(0).squeeze().cuda()
emb_mask,_ = mask.max(1)
这里值得注意的是mask的设置。作者将被分配为目标（targets）的anchor设置mask，那么只考虑被分配为targets的anchor对应位置上的值来计算损失函数。
（4）embedding损失的计算
接着作者计算embedding损失，值得注意的是，在该过程中作者采用了前面提到的全连结层来获得embedding的高级语义信息（track ID）。然后使用常见用于分类任务的交叉熵损失函数。
具体代码如下：
# For convenience we use max(1) to decide the id, TODO: more reseanable strategy
tids,_ = tids.max(1)
tids = tids[emb_mask]
embedding = p_emb[emb_mask].contiguous()
embedding = self.emb_scale * F.normalize(embedding)
nI = emb_mask.sum().float()
if test_emb:
if np.prod(embedding.shape)==0 or np.prod(tids.shape) == 0:
return torch.zeros(0, self.emb_dim+1).cuda()
emb_and_gt = torch.cat([embedding, tids.float()], dim=1)
return emb_and_gt
if len(embedding) > 1:
logits = classifier(embedding).contiguous()
lid = self.IDLoss(logits, tids.squeeze())
（5）总损失函数
最后，作者按照公式
对上面的三个损失进行求和。实现代码如下
loss = torch.exp(-self.s_r)*lbox + torch.exp(-self.s_c)*lconf + torch.exp(-self.s_id)*lid + \
(self.s_r + self.s_c + self.s_id)
loss *= 0.5
至此，有关predicition head的部分就讲解结束了。JDE的主体部分就介绍完了，其他细节，大家可以看一下原论文和代码，进行探索。
3 总结
tracking by detection是非常常见的MOT范式，但是目前MOT领域为了平衡追踪速度和精度，慢慢放弃了这种范式，转而投入将检测与embedding/匹配进行结合的范式研究中。本文介绍的JDE就是一个网络同时输出图像画面中的检测框位置和检测框内物体的embedding，从而加速MOT的速度。
但是值得注意的是，JDE只是同时输出了检测框和embedding信息。后面还是要通过卡尔曼滤波和匈牙利算法进行目标的匹配。总的来说，还是分为检测和匹配双阶段。
