源码解析目标检测的跨界之星DETR（五）、loss函数与匈牙利匹

2020-07-17 本文已影响0人 CW不要无聊的风格

Date: 2020/07/17

Coder: CW

Foreword:

本文将对 loss函数的实现进行解析，由于 DETR 是预测结果是集合的形式，因此在计算loss的时候有个关键的前置步骤就是将预测结果和GT进行匹配，这里的GT类别是不包括背景的，未被匹配的预测结果就自动被归类为背景。匹配使用的是匈牙利算法，该算法主要用于解决与二分图匹配相关的问题，对这部分感兴趣的朋友们可以参考下这篇文：匈牙利算法

Outline

I. Loss Function

i). 分类loss

ii). 回归loss

II. Hungarian algorithm（匈牙利算法）

Loss Function

先来看看与loss函数相关的一些参数：matcher就是将预测结果与GT进行匹配的匈牙利算法，这部分的实现会在下一节解析。weight_dict是为各部分loss设置的权重，主要包括分类与回归损失，分类使用的是交叉熵损失，而回归损失包括bbox的 L1 Loss（计算x、y、w、h的绝对值误差）与 GIoU Loss。若设置了masks参数，则代表分割任务，那么还需加入对应的loss类型。另外，若设置了aux_loss，即代表需要计算解码器中间层预测结果对应的loss，那么也要设置对应的loss权重。

与loss函数实现相关的初始化参数

loss函数是通过实例化SetCriterion对象来构建。

构建loss函数

losses变量指示需要计算哪些类型的loss，其中cardinality仅用作log，并不涉及反向传播梯度。

loss_cardinality

可以先来看下SetCriterion这个类的doc string，了解下各部分参数的意义。

SetCriterion(i)

CW 也作了对应的注释：

SetCriterion(ii)

接下来看下其前向过程，从而知悉loss的计算。

这里一定要先搞清楚模型的输出（outputs）和GT（targets）的形式，对于outputs可参考CW在下图中的注释；而targets是一个包含多个dict的list，长度与batch size相等，其中每个dict的形式如COCO数据集的标注，具体可参考该系列的第二篇文章： 源码解析目标检测的跨界之星DETR（二）、模型训练过程与数据处理 中的数据处理部分。

SetCriterion(iii)

如CW在前言部分所述，计算loss的一个关键的前置步骤就是将模型输出的预测结果与GT进行匹配，对应下图中self.matcher()的部分，返回的indices的形式已在注释中说明。

SetCriterion(iv)

接下来是计算各种类型的loss，并将对应结果存到一个dict中（如下图losses变量），self.get_loss()方法返回loss计算结果。

SetCriterion(v)

SetCriterion(vi)

get_loss方法中并不涉及具体loss的计算，其仅仅是将不同类型的loss计算映射到对应的方法，最后将计算结果返回。

get_loss

接下来，我们就对分类和回归损失的计算过程分别进行解析。

i). 分类loss

首先说明下，doc string里写的是NLL Loss，但实际调用的是CE Loss，这是因为在Pytorch实现中，CE Loss实质上就是将Log-Softmax操作和NLL Loss封装在了一起，如果直接使用NLL Loss，那么需要先对预测结果作Log-Softmax操作，而使用CELoss则直接免去了这一步。

loss_labels(i)

其次，要理解红框部分的_get_src_permutation_idx()在做什么。输入参数indices是匹配的预测（query）索引与GT的索引，其形式在上述SetCriterion(iv)图中注释已有说明。该方法返回一个tuple，代表所有匹配的预测结果的batch index（在当前batch中属于第几张图像）和 query index（图像中的第几个query对象）。

_get_src_permutation_idx

类似地，我们可以获得当前batch中所有匹配的GT所属的类别（target_classes_o），然后通过src_logits、target_classes_o就可以设置预测结果对应的GT了，这就是下图中的target_classes。target_classes的shape和src_logits一致，代表每个query objects对应的GT，首先将它们全部初始化为背景，然后根据匹配的索引（idx）设置匹配的GT（target_classes_p）类别。

loss_labels(ii)

“热身活动”做完后，终于可以开始计算loss了，注意在使用Pytorch的交叉熵损失时，需要将预测类别的那个维度转换到通道这个维度上（dim1）。

loss_labels(iii)

另外，class_error计算的是Top-1精度（百分数），即预测概率最大的那个类别与对应被分配的GT类别是否一致，这部分仅用于log，并不参与模型训练。

accuracy

ii). 回归loss

回归loss的计算包括预测框与GT的中心点和宽高的L1 loss以及GIoU loss。

注意在下图注释中，num_matched_queries1+num_matched_queries2+..., 和 num_matched_objs1+num_matched_objs2+... 是相等的，在前面 SetCriterion(iv) 图中matcher的返回结果注释中有说明。

loss_boxes(i)

以下就是loss的计算。注意下 reduction 参数，若不显式进行设置，在Pytorch的实现中默认是'mean'，即返回所有涉及误差计算的元素的均值。

loss_boxes(ii)

另外，在计算GIoU loss时，使用了torch.diag()获取对角线元素，这是因为generalized_box_iou()方法返回的是所有预测框与所有GT的GIoU，比如预测框有N个，GT有M个，那么返回结果就是NxM个GIoU。而如 loss_boxes(i) 图中所示，我们预先对匹配的预测框和GT进行了排列，即N个预测框中的第1个匹配M个GT中的第1个，N中第2个匹配M中第2个，..，N中第i个匹配M中第i个，于是我们要取相互匹配的那一项来计算loss。

generalized_box_iou(i)

generalized_box_iou(ii)

Hungarian algorithm（匈牙利算法）

build_matcher()方法返回HungarianMatcher对象，其实现了匈牙利算法，在这里用于预测集（prediction set）和GT的匹配，最终匹配方案是选取“loss总和”最小的分配方式。注意CW对loss总和这几个字用了引号，其与loss函数中计算的loss并不完全一致，在这里是作为度量（cost/metric）的角色，度量的值决定了匹配的结果，接下来我们看代码实现就会一清二楚。