用 Keras/TensorFlow 2.8 创建 COCO 的 average precision 指标

---

前言

YOLO 系列（包括 YOLOv4-CSP，YOLOv4 等）的探测器 detector，它们的损失函数由 3* 部分组成。如果这 3 部分的每一个部分都使用一个指标，将得到 3 个独立的指标。这会使得训练模型变得更困难（比如一个指标变好，另一个指标却变坏的情况）。
正如吴恩达教授在《机器学习策略》（Introduction to Machine Learning Strategy）课程中所提到的，当一个模型有多个指标时，应该尽量把它们合并成一个最重要的优化指标，才能使得模型有明确的优化方向。而 COCO 数据集的 average precision 指标，就是最合适的优化指标。

• *这 3 部分损失是：1. 判断预设框 anchor box 内是否有物体的预测损失。2. 判断预设框内的物体，属于哪一个类别的预测损失。3. 预测结果物体框和标签物体框，两者之间的 CIOU 损失。

---

1. AP 的算法原理。

COCO 数据集的 AP(average precision) 指标，实际上是对 average precision 计算了两次平均值，主要算法有如下 3 个步骤：

1. 设定 IoU 阈值为 0.5，计算一个类别的 AP。
2. 重复上面的第一步，计算 80 个类别 AP，然后对 80 个类别 AP 求平均值，得到一个 average_precision_over_categories。
3. 遍历 10 个 IoU 阈值（0.5，0.55， 0.6，…， 0.95，以 0.05 为步进值），重复上面的 2 个步骤，得到 10 个 average_precision_over_categories，然后对这 10 个 average_precision_over_categories 求平均值，就得到一个 mean average precision。这个 mean average precision 就是 COCO 的 AP。

COCO 数据集声明不区分 AP 和 mAP（mean average precision），统一使用 AP 这个词，由使用者根据使用场景自行区分是 AP 还是 mAP。如果想要区分的话，简单来说，对单个类别就是 AP，求了两次平均值之后就是 mAP（mean average precision）。

---

2. 在 Keras 中的实现。

下面我们用 Keras/TensorFlow 2.8 来创建 COCO 的 AP 指标，并详细解释其算法原理。

Keras 中有一个专门用于指标的类： tf.keras.metrics.Metric。可以用它来创建一个类 MeanAveragePrecision，来计算 AP 指标。

指标类 tf.keras.metrics.Metric 中，有 3 个主要的方法，update_state、 result 和 reset_state。这 3 个方法的具体作用是：

1. 方法 update_state，根据每个 batch 的计算结果，对状态量进行更新。
2. 方法 result，使用更新好的状态量，计算指标。
3. 方法 reset_state，用于在每个 epoch 开始时，把状态量重新设置为初始状态。

在代码中的指标 MeanAveragePrecision 和 3 个方法如下图。

对于 COCO 的 AP 来说，主要用到类别置信度和 IoU 这两个数据，所以方法 update_state 将主要更新这 2 个数据，而方法 result 也将主要使用这 2 个数据来计算 AP。

下面的程序伪代码中，会经常提到 4 个词语，这里先明确一下这 4 个词语的定义：

1. objectness：是一个概率值，表示对于当前物体框，框内有物体存在的概率。（YOLO 论文中使用该词，中文的翻译应该是 “物体框内有物体存在的置信度”。为了便于和下面第二条的类别置信度进行区别，这里沿用论文的写法，用 objectness。）
2. 类别置信度：也是一个概率值，表示对于当前物体框内物体，属于某个类别的概率大小。
3. “正样本”：正样本的意思，是指一个 bbox（它的 objectness 和类别置信度都大于对应的阈值）。
4. 相关图片：是指该图片的标签或是预测结果的正样本中，包含了该类别。

---

3. 创建状态量。

用指标类 tf.keras.metrics.Metric 创建的是完整状态的指标 stateful metric，意思是可以用它来计算很复杂的指标。通常需要先创建相关的状态量 states。

对于 COCO 的 AP 指标，需要创建 3 个状态量 latest_positive_bboxes、 labels_quantity_per_image 和 showed_up_classes，3 个状态量类型均为 tf.Variable。

1. latest_positive_bboxes，形状为 (CLASSES, latest_related_images, bboxes_per_image, 2)。记录的是对于每一个类别，使用 latest_related_images 张相关图片，计算得到的状态值。
   而对于每一张相关图片，其对应的状态张量形状为 (bboxes_per_image, 2)，记录了 bboxes_per_image 个 bboxes 的状态。
   每个 bboxes 的状态是一个长度为 2 的向量，两个值分别是类别置信度和 IoU 值。
2. labels_quantity_per_image，形状为 (CLASSES, latest_related_images)。记录的是对于每一个类别，在 latest_related_images 张相关图片中，标签 bboxes 的数量。
   这个状态张量和 latest_positive_bboxes 是一一对应的。也就是说，如果状态 latest_positive_bboxes 中记录了一个相关图片的置信度和 IoU，则 labels_quantity_per_image 也必须记录该图片中的标签 bboxes 数量。
3. showed_up_classes，形状为 (CLASSES, )，是一个布尔张量，记录的是所有图片中出现过的类别。
   在计算 AP 时，只有出现过的类别，才能参与计算 AP，所以需要用 showed_up_classes 进行记录。

创建好的 3 个状态量如下图。

在上图中有 3 点要注意：

1. 为了实现 P3, P4, P5 共用状态量 states，把状态量建立在类 MeanAveragePrecision 的外部。这是因为 YOLOv4-CSP 和 YOLOv4 等模型，有 P3, P4, P5 一共 3 个输出。如果状态量建立在 MeanAveragePrecision 的内部，这些状态量实际上将会是复制了 3 份，即 3x3=9 个独立的状态量。
2. 设置 trainable=False。因为这 3 个状态量是 tf.Variable，默认会求梯度并进行反向传播。但是对指标的状态量来说，并不需要进行反向传播，所以设置 trainable=False，可以节省计算资源。
3. 状态量的形状，受 2 个全局变量 latest_related_images, bboxes_per_image 控制。如果电脑的算力足够，可以把这两个变量设置得大一些。

---

4. update_state 方法。

在每个 batch 计算完成之后，要用方法 update_state 对 3 个状态量进行更新。

4.1 更新第一个状态量 showed_up_classes 。

下面是程序的伪代码，用到一些变量的名字，和程序中的变量名字相同，以方便阅读代码。

1. 从标签中提取出现过的类别，得到张量 showed_up_categories_label，张量形状为 (x，),里面存放的是出现过的类别编号，表示有 x 个类别出现在了这批标签中。
   1.1 showed_up_categories_index_label = tf.experimental.numpy.isclose(objectness_label, 1)
   1.2 showed_up_categories_label = tf.argmax(y_true[…, 1: 81], axis=-1)，showed_up_categories_label 形状为 (batch_size, *Feature_Map_px, 3)。在代码中会以 *Feature_Map_px 表示 P5, P4, P3 的特征图大小，一般 P5 特征图大小为（19, 19）。
   1.3 showed_up_categories_label = showed_up_categories_label[showed_up_categories_index_label]，showed_up_categories_label 形状为 (x, )。

   注意不能仅使用 argmax，而是必须借助上面的第 1.3 步骤，使用 showed_up_categories_index_label 作为索引。因为在没有标签时，argmax 也会得出一个值 0，而这个 0 并不表示该物体框的类别为 0 。

2. 从预测结果中提取出现过的类别，得到张量 showed_up_categories_pred，张量形状为 (y，),里面存放的是出现过的类别编号，表示有 y 个类别出现在了这批预测结果中。

3. 将上面 2 个张量改变形状为 (1, -1)，然后用 tf.sets.union 求并集，得到一个 sparse tensor， 将其转换为 tf.tensor，得到张量 showed_up_categories_batch，张量形状为 (categories_batch,)。

4. 遍历 showed_up_categories_batch，对每一个出现过的类别 category，如果之前还没有出现过，则设置 showed_up_classes[category].assign(True)。

4.2 更新另外两个状态量。

3 大主要操作步骤如下：

1. 第一重循环，遍历批次结果数据中的每一张图片（方法 update_state 中接收的是单个批次的结果，所以要遍历批次中的每张图片），对每一张图片执行下面 2 步操作。

2. 对于单张图片的标签 one_label 和预测结果 one_pred，分别构造相应的张量。

   2.1 对于标签：构造张量 positives_index_label，positives_label 和 category_label。

   positives_index_label 形状为 (19, 19, 3)，是一个布尔张量，是标签正样本的索引张量，即只有 objectness 等于 1 的 bboxes，其对应布尔值才会为 True。
   为了方便描述，这里的形状只以 YOLOv4-CSP 模型的 P5 形状为例。下面也是如此。
   创建标签的正样本张量：
   positives_label = tf.where(condition=positives_index_label[…, tf.newaxis], x=one_label, y=-8.0]，positives_label 形状为 (19, 19, 3，85), 里面只有正样本信息，其它位置的数值为 -8。（使用 -8 而没有使用 -1，是为了便于和 axis=-1 混淆，方便搜索）

   创建标签的类别张量：
   category_label = tf.math.argmax(positives_label[…, 1: 81], axis=-1)，category_label 形状为 (19, 19, 3)，代表每一个正样本的类别。在不是正样本的位置，数值为 0。因为这个 0 会和类别编号 0 发生混淆，所以下面要用 tf.where 再次进行转换，使得在不是正样本的位置，其数值为 -8。
   category_label = tf.where(condition=positives_index_label, x=category_label, y=-8)

   2.2 对于预测结果：构造张量 positives_index_pred，positives_pred 和 category_pred。

   positives_index_pred 形状为 (19, 19, 3)，是一个布尔张量，是预测结果正样本的索引张量，即只有 objectness 和类别置信度都大于阈值的 bboxes，其对应布尔值才会为 True。
   构造另外 2 个张量的方法，和构造对应的标签张量方法相同。这里不再赘述。

3. 第二重循环，遍历 80 个类别，区分 4 种情况，更新状态值。

先创建标签类别的布尔值 category_bool_any_label，和预测结果的布尔值 category_bool_any_pred，两者均为标量型张量，如果有任何一个物体框内物体属于当前类别，则对应布尔值为 True：

category_bool_label = tf.experimental.numpy.is_close(category_label, category)
category_bool_any_label = tf.reduce_any(category_bool_label)

category_bool_pred = tf.experimental.numpy.is_close(category_pred, category)
category_bool_any_pred = tf.reduce_any(category_bool_pred)

对每一个类别，都要区分 4 种情况，计算得到两个张量 one_image_positive_bboxes 和 one_image_category_labels_quantity，分别用来更新两个状态量 latest_positive_bboxes 和 labels_quantity_per_image。

对 4 种情况构建布尔张量：
情况 a ：标签和预测结果中，都没有该类别。无须更新状态。
情况 b ：预测结果中没有该类别，但是标签中有该类别。布尔张量为 scenario_b = tf.logical_and(~category_bool_any_pred, category_bool_any_label)。
此时需要提取预测结果的类别置信度和 IoU，且类别置信度和 IoU 都为 0。另外还需要提取标签数量。

情况 c ：预测结果中有该类别，标签没有该类别。布尔张量为 scenario_c = tf.logical_and(category_bool_any_pred, ~category_bool_any_label)。
此时需要提取预测结果的类别置信度和 IoU。而因为没有标签，IoU 为0。另外还需要提取标签数量 0。

情况 d ：预测结果和标签中都有该类别，布尔张量为 scenario_d = tf.logical_and(category_bool_any_pred, category_bool_any_label)。
此时需要提取预测结果的类别置信度和 IoU。IoU 要经过计算得到。另外还需要提取标签数量。

只有在情况 b,c,d 时，才需要更新 2 个状态量，所以先要判断是否处在情况 b,c,d 下，再决定是否执行后续步骤。

under_scenarios_bc = tf.logical_or(scenario_b, scenario_c)
under_scenarios_bcd = tf.logical_or(under_scenarios_bc, scenario_d)

if under_scenarios_bcd:
    提取 b,c,d 三种情况的置信度和 IoU，更新另外 2 个状态量。

对于每一个类别来说，a,b,c,d 情况不会同时发生，只可能出现其中的一种，因为这 4 种情况是互斥的。
下面是 b,c,d 情况下，详细的操作步骤。

情况 b：
one_image_positive_bboxes = tf.zeros(shape=(BBOXES_PER_IMAGE, 2))，可以直接把 one_image_positive_bboxes 作为输出张量。

情况 c，需要提取类别置信度和 IoU，有 5 个操作步骤：

1. 先获取当前情况的正样本：
   scenario_c_positives_pred = positives_pred[category_bool_pred]，scenario_c_positives_pred 形状为 (scenario_c_bboxes, 85)。

2. 再获取当前情况的类别置信度：
   scenario_c_classification_confidence_pred = tf.reduce_max(scenario_c_positives_pred[:, 1: 81])，scenario_c_positives_pred 形状为 (scenario_c_bboxes,)。

3. 比较 scenario_c_bboxes 和 bboxes_per_image 的大小，区分两种情况：

   3.1 如果 scenario_c_bboxes ＜ bboxes_per_image：
   使用 tf.pad，对 scenario_c_classification_confidence_pred 尾部进行补零，得到新的张量 one_image_positive_bboxes，其形状变为 (bboxes_per_image,) 。

   3.2 如果 scenario_c_bboxes ≥ bboxes_per_image：

   按照置信度从大到小的顺序，对 scenario_c_classification_confidence_pred 进行排序，得到 scenario_c_sorted_pred，其形状为 (scenario_c_bboxes,)。
   之所以要进行排序，是因为在最后计算 AP 时，需要按置信度从大到小进行排序后，才会计算 AP。所以当前步骤在筛选 bboxes 时，也就应该把置信度大的 bboxes 筛选出来。

   保留置信度较大的 bboxes，即 one_image_positive_bboxes = scenario_c_sorted_pred[:bboxes_per_image]，one_image_positive_bboxes 形状为 (bboxes_per_image,) 。

4. 获取 IoU（情况 c 的 IoU 为 0，情况 d 的 IoU 需要经过计算得到）。
   因为标签中没有这个类别，所以 IoU 为 0，可以使用 scenario_c_ious_pred = tf.zeros_like(one_image_positive_bboxes).

5. 将置信度和 IoU 进行堆叠 stack。
   one_image_positive_bboxes = tf.stack(values=[one_image_positive_bboxes, scenario_c_ious_pred], axis=1)，one_image_positive_bboxes 形状为 (bboxes_per_image, 2)。

情况 d：此时需要计算 IoU，6 个步骤如下（因为只有和标签 IoU 最大的预测结果 bbox，才认为是命中了该标签，所以需要对每一个标签，同时和所有的预测结果 bboxes 计算 IoU）：

1. 对于预测结果：取出属于当前类别的 bboxes，把每个 bbox 信息填入全零数组 bboxes_iou_pred，其它多余的位置保持数值为 0。bboxes_iou_pred = tf.where(condition=category_bool_pred, x=positives_pred[…, -4:], y=0]，bboxes_iou_pred 形状为 (19, 19, 3，4)。

2. 对于标签：取出属于当前类别的 bboxes，即 bboxes_category_label = positives_label[…, -4:][category_bool_label]，bboxes_category_label 形状为 (scenario_d_bboxes_label，4)。

3. 对 bboxes_category_label，按照面积从小到大的顺序进行排序（体现着重小物体的思想，善于识别小物体的模型，其指标将越好），得到 sorted_bboxes_label，形状为 (scenario_d_bboxes_label, 4)。

4. 建立张量 one_image_positive_bboxes = tf.zeros(shape=(BBOXES_PER_IMAGE, 2))。设置计数变量 new_bboxes_quantity = 0.

5. 遍历 sorted_bboxes_label，对每一个标签 bbox，执行如下 3 个操作：

   5.1 把其信息填入全 1 张量 bbox_iou_label（即张量最后一个维度，所有长度为 4 的向量，写的都是同一个 bbox 的信息），bboxes_iou_pred 形状为 (19, 19, 3, 4)。

   5.2 用 bbox_iou_label 和 bboxes_iou_pred 计算 ious_category，ious_category 张量形状为 (19, 19, 3)。

   5.3 如果最大 IoU 大于阈值 0.5，则认为预测结果中对应的 bbox 命中了标签，做 2 个操作：
   5.3.1 将该 bbox 的类别置信度和 IoU 记录到张量 one_image_positive_bboxes 中（用 tf.concat）。
   5.3.2 从 bboxes_iou_pred 去掉该 bbox(用 tf.where)，后续计算 IoU 不需要再考虑这个 bbox。
   5.3.3 new_bboxes_quantity += 1.

   5.4 new_bboxes_quantity 等于 bboxes_per_image 时，停止记录新的 bboxes。

6. 遍历 sorted_bboxes_label 完成之后，如果 bboxes_iou_pred 有剩余的 bboxes，说明这些 bboxes 没有命中任何标签。如果还满足条件 new_bboxes_quantity < BBOXES_PER_IMAGE，则需要将剩下 bboxes 的 IoU 设为 0，并记录到张量 one_image_positive_bboxes 中。

   令剩余 bboxes 数量为 left_bboxes_quantity, 加到 one_image_positive_bboxes 后，总的 bboxes 数量为 scenario_d_bboxes = new_bboxes_quantity + left_bboxes_quantity。

   求出剩余的 bboxes，得到 left_bboxes_pred， 形状为 (left_bboxes_quantity, 85)。
   left_bboxes_confidence_pred = tf.reduce_max(left_bboxes_pred[:, 1: 81])

   6.1 如果 scenario_d_bboxes ＞ bboxes_per_image，则需要进行排序：

   6.1.1 按照置信度从大到小的顺序，对 left_bboxes_confidence_pred 进行排序，得到 left_bboxes_sorted_confidence，其形状为 (left_bboxes_quantity,)。
   之所以要进行排序，是因为在最后计算 AP 时，需要按置信度从大到小进行排序后，才会计算 AP。
   6.1.2 vacant_seats = BBOXES_PER_IMAGE - new_bboxes_quantity
   6.1.3 left_bboxes_confidence_pred = left_bboxes_sorted_confidence[:vacant_seats]。

   6.2 如果 scenario_d_bboxes ≤ bboxes_per_image，则无须进行排序，可以直接使用left_bboxes_confidence_pred。

   6.3 给 left_bboxes_confidence_pred 加上全为 0 的 IoU (使用 tf.stack)，得到 left_positive_bboxes_pred，形状为(vacant_seats, 2)。

   6.4 将 left_positive_bboxes_pred 和 one_image_positive_bboxes 进行拼接 concatenate，然后保留最后 bboxes_per_image 个 bboxes，即 one_image_positive_bboxes = one_image_positive_bboxes[-bboxes_per_image：]。

计算标签数量，得到整数 one_image_category_labels_quantity = tf.where(category_bool_label).shape[0]。

最后更新 2 个状态量，更新原则为先进先出 FIFO。

1. 用张量 one_image_positive_bboxes 更新状态量 latest_positive_bboxes。
latest_positive_bboxes 形状为 (CLASSES, latest_related_images, bboxes_per_image, 2)。

latest_positive_bboxes[category, 1:].assign(latest_positive_bboxes[category, :-1])
latest_positive_bboxes[category, 0].assign(one_image_positive_bboxes)


2. 用整数 one_image_category_labels_quantity 更新状态量 labels_quantity_per_image。
labels_quantity_per_image 形状为 (CLASSES, latest_related_images)。

labels_quantity_per_image[category, 1:].assign(labels_quantity_per_image[category, :-1])
labels_quantity_per_image[category, 0].assign(one_image_category_labels_quantity)

---

5. result 方法。

方法 result 的作用，是使用状态量来计算指标。

需要注意的是，YOLO-v4-CSP 是多输出模型，有 P3, P4, P5 这 3 个输出，所以在每批次数据计算完成之后，会在这 3 个输出上分别计算一次指标。可以设置跳过 P4, P5，只计算 P3 的指标。

在方法 result 中，根据自顶向下的程序结构，顶层的程序只有 2 个大步骤：

1. 遍历 10 个 IoU 阈值，对每一个 IoU 阈值，计算 1 个 average_precision_over_categories：
   1.1 遍历 80 个类别，对每一个类别，计算 1 个 AP。
   1.2 对 80 个 AP 取平均值，得到 average_precision_over_categories。
2. 将最终的 10 个 average_precision_over_categories 取平均值，就得到最终的 mAP。

在上面的步骤 1.1 中，计算单个类别的 average_precision 时，如果 labels_quantity = 0，直接设 AP = 0。

而如果 labels_quantity 不等于 0，则需要计算 AP，有如下 2 个操作：

1. 计算 recall_precisions。
   1.1 创建空的张量 recall_precisions，形状为 (1,)。其索引为 recall，设置初始 recall = 0, recall_precisions[0] = 1。后续每一个 recall 值都对应一个 precision 值。
   1.2 设置 true_positives = 0，false_positives = 0。
   1.3 从 latest_positive_bboxes 中，取出当前类别的所有 bboxes，形状是 (bboxes_per_image * latest_related_images, 2)。每个 bbox 包含 2 个信息：类别置信度和 IoU。
   1.4 按照类别置信度，进行由大到小的排序，得到张量 sorted_bboxes_category。

   1.5 遍历 sorted_bboxes_category 中的所有 bboxes，计算得到 recall_precisions (使用类别置信度进行过滤。如果置信度为 0，说明它不是模型的预测结果，不应该参与计算 AP)。
   1.5.1 如果该 bbox 的 IoU 大于当前的 IoU 阈值，则认为该预测正样本命中，更新 2 个数值， true_positives += 1， recall += 1。
   1.5.2 如果 IoU 小于阈值，则更新 1 个数值 false_positives += 1。

   1.5.3 计算 precision = true_positives/（false_positives + true_positives），然后更新张量 recall_precisions，即 recall_precisions[recall] = precision 。

2. 计算 AP。
   遍历 sorted_bboxes_category 完成后，使用张量 recall_precisions，计算多个小梯形面积，累加所有小梯形的面积，得到 AP。

   2.1 从 labels_quantity_per_image 中，获得 latest_images 个相关图片的标签 bboxes 总数 labels_quantity，而 1/labels_quantity 则是小梯形的高度 trapezoid_height。

   2.2 从 recall_precisions 的第 0 个索引位置开始，计算小梯形面积 (recall_precisions[0] + recall_precisions[1]) * trapezoid_height / 2。直到 recall_precisions 的倒数第 2 个索引位置结束。

   2.3 把所有的小梯形面积累加起来，就得到该类别的 AP。

---

6. 测试盒 testcase。

做了一个测试盒 testcase，盒子里放了 13 个单元测试。目前 AP 指标通过了盒子里全部的 13 个测试。

如果使用者需要改动这个指标，也应该用测试盒再测试一下，确保指标能正常运行。

对于企业用户，当然应该按照软件工程的要求，由软件测试团队进行专业的测试之后，才能使用。
测试盒程序的部分截图如下：

---

7. 使用方法。

在使用这个 AP 指标文件时，注意以下 3 点：

1. 该指标需要配合使用 YOLO 系列的模型，比如 YOLOv4-CSP, YOLOv4 等等。
   因为 YOLO 系列的模型有 p3, p4, p5 共 3 个输出，指标也是针对这 3 个输出写的。如果需要用到其它的探测器 detector上，需要自行修改指标文件。
2. 该指标文件可以运行在 TensorFlow 2.8 环境下。如果要使用更低版本的 TensorFlow，可能需要自行修改文件中的少量代码。
   举例来说，如下左图，在 TF 2.4 中，isclose 函数的结果是一种特殊的 TF 数组。该数组无法直接用做张量的索引，需要先手动将其转换成张量。这可以算作 TF 2.4 的一个 bug。
   而到了 TF 2.8，修复了这个问题，isclose 函数的结果是一个 TF 张量，可以直接用做张量的索引，方便了很多。如下右图。
3. 该指标可能需要在 eager 模式下运行。
   在图模式下，该指标会生成计算图，将占用大量的内存，超过 128G，所以个人的台式机难以将其运行在图模式下，需要使用 eager 模式。
   对于企业用户，有大量的内存和算力的条件下，可以尝试用图模式。
   要在 eager 模式下运行该指标，直接在编译模型时设置 run_eagerly=True 即可，示例如下：

yolo_v4_csp_model.compile(
	run_eagerly=True, 
	metrics=average_precision,
	loss=my_custom_loss, 
	optimizer=optimizer_adam)

---

8. 下载链接。

代码已在 Github 开源，可以直接下载。→ 下载链接在此
一共有 2 个相关文件，指标文件 average_precision_metric.py 和测试盒文件 testcase_average_precision.py。

---

THE END