MMSegmentation和Detectron2中几种Metrics的实现

发表于 更新于 分类于 阅读次数：

本文记录MMSegmentation和Detectron2中Metircs模块的实现，关注模块设计，实现细节，边界情况，给出在库外使用该模块的方法；MMSegmentation中这部分的代码主要在此处，Detectron2中这部分的代码主要在此处。

常用Metrics定义

下列定义来自FCN。其中，\(K\) 是总类别数，\(i\) 是类别索引，\(n_{ij}\) 表示真实类别为 \(i\) 预测类别为 \(j\) 的总像素数，\(t_i = \sum_{j}n_{ij}\) 是第 \(i\) 类别所有像素数量。

pAcc

一个像素视为一个统计单位，计算准确率。

\[ \mathrm{pAcc} = \frac{\sum_{i}n_{ii}}{\sum_{i}t_i} \]

mAcc

mIoU

fwIoU

模块设计

1
2
3
4
eval_metrics
├─total_intersect_and_union
│  └─intersect_and_union
├─total_area_to_metrics
  • total_intersect_and_union 循环为多个(pred, gt)调用 intersect_and_union，并使用加法在数量维度上规约其结果；
  • intersect_and_union 计算四个的直方图（横坐标为类，纵坐标为像素数量）：交、并、pred，gt；
  • total_area_to_metrics 使用总直方图信息计算
  • 诸如 mean_iou, mean_dice, mean_fscore都是对eval_metrics 的简单封装（这三个也是目前MMSeg支持的，除了Acc, Precision, Recall之外的指标）；
  • pre_eval_to_metrics 将 pre_eval 积累的所有输出的评估结果加法规约，之后再送入 total_area_to_metrics；

什么情况下触发 pre_eval？为什么需要 pre_eval？

  • 在使用 --eval 指定评估指标，并且不是在cityscapes的format结果上进行cityscape mIoU的评估时，都会开启 pre_eval；
  • 希望在将验证集上每个batch的输入（事实上在验证/测试时batch=1）通过模型后，就立刻与gt计算metrics，而不是在某处保存生成的结果，这样将来在reduce的时候也无需再次加载gt；
  • 按照笔者曾经的经验来看，由于把每个batch的结果都添加到内存里，笔者曾经用一台内存比较小的机器跑ADE20K的推理，观察到内存占用逐渐增加，最后出现内存的OOM问题；为什么不边算便求和？

实现细节

这里主要关注 intersect_and_union 以及 total_area_to_metrics 的实现；

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def intersect_and_union(pred_label,
                        label,
                        num_classes,
                        ignore_index,
                        label_map=dict(),
                        reduce_zero_label=False):

    if isinstance(pred_label, str):
        pred_label = torch.from_numpy(np.load(pred_label))
    else:
        pred_label = torch.from_numpy((pred_label))

    if isinstance(label, str):
        label = torch.from_numpy(
            mmcv.imread(label, flag='unchanged', backend='pillow'))
    else:
        label = torch.from_numpy(label)

    if reduce_zero_label:
        label[label == 0] = 255
        label = label - 1
        label[label == 254] = 255
    if label_map is not None:
        label_copy = label.clone()
        for old_id, new_id in label_map.items():
            label[label_copy == old_id] = new_id

    mask = (label != ignore_index)
    pred_label = pred_label[mask]
    label = label[mask]

    intersect = pred_label[pred_label == label]
    area_intersect = torch.histc(
        intersect.float(), bins=(num_classes), min=0, max=num_classes - 1)
    area_pred_label = torch.histc(
        pred_label.float(), bins=(num_classes), min=0, max=num_classes - 1)
    area_label = torch.histc(
        label.float(), bins=(num_classes), min=0, max=num_classes - 1)
    area_union = area_pred_label + area_label - area_intersect
    return area_intersect, area_union, area_pred_label, area_label


def total_area_to_metrics(total_area_intersect,
                          total_area_union,
                          total_area_pred_label,
                          total_area_label,
                          metrics=['mIoU'],
                          nan_to_num=None,
                          beta=1):

    if isinstance(metrics, str):
        metrics = [metrics]
    allowed_metrics = ['mIoU', 'mDice', 'mFscore']
    if not set(metrics).issubset(set(allowed_metrics)):
        raise KeyError('metrics {} is not supported'.format(metrics))

    all_acc = total_area_intersect.sum() / total_area_label.sum()
    ret_metrics = OrderedDict({'aAcc': all_acc})
    for metric in metrics:
        if metric == 'mIoU':
            iou = total_area_intersect / total_area_union
            acc = total_area_intersect / total_area_label
            ret_metrics['IoU'] = iou
            ret_metrics['Acc'] = acc
        elif metric == 'mDice':
            dice = 2 * total_area_intersect / (
                total_area_pred_label + total_area_label)
            acc = total_area_intersect / total_area_label
            ret_metrics['Dice'] = dice
            ret_metrics['Acc'] = acc
        elif metric == 'mFscore':
            precision = total_area_intersect / total_area_pred_label
            recall = total_area_intersect / total_area_label
            f_value = torch.tensor(
                [f_score(x[0], x[1], beta) for x in zip(precision, recall)])
            ret_metrics['Fscore'] = f_value
            ret_metrics['Precision'] = precision
            ret_metrics['Recall'] = recall

    ret_metrics = {
        metric: value.numpy()
        for metric, value in ret_metrics.items()
    }
    if nan_to_num is not None:
        ret_metrics = OrderedDict({
            metric: np.nan_to_num(metric_value, nan=nan_to_num)
            for metric, metric_value in ret_metrics.items()
        })
    return ret_metrics

在 intersect_and_union 使用了 torch.histc 实现，避免了逐个类的for循环；

值得注意的是，所有的metrics都是先对直方图信息求和，再计算，当然求和的位置确实会影响最终结果，笔者认为这样做可以避免频繁在个例中出现Nan；以第 \(c\) 类的 \(\mathrm{IoU}_c\) 为例：

\[ \mathrm{IoU}_c =\frac{ \sum_{i=1}^N P_c^{(i)}\cap Q_c^{(i)}}{ \sum_{i=1}^N P_c^{(i)}\cup Q_c^{(i)}} \]

边界情况：当某个类别从未在任何一张gt上出现时，以 total_area_label 为分母的metrics都是Nan，如果还满足从未在任何一张pred上出现，那么 total_area_union 为分母的IoU也是Nan；

在任何地方使用

笔者认为，使用 mean_xxx 这类接口就能满足大部分的需求，以下面一个随机demo为例，这个例子也验证了IoU的计算并不是逐对计算后平均；

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
from mmseg.core.evaluation import mean_iou
import numpy as np
from functools import partial

def random_int_matrix(shape, upper):
    tmp = np.zeros(shape)
    tmp.fill(upper)
    return np.random.randint(tmp)


def test_order():
    gen_mat = partial(random_int_matrix, (10, 10), 5)
    cal_mean_iou = partial(mean_iou, num_classes=5, ignore_index=-1)
    N = 100
    ls = [gen_mat() for i in range(N)]
    ps = [gen_mat() for i in range(N)]
    miou = cal_mean_iou(ps, ls)

    ious = []
    for p, l in zip(ps, ls):
        ious.append(cal_mean_iou(p, l)['IoU'])

    avg_iou = sum(ious) / len(ious)
    print(f'miou: {miou}')
    print(f'avg_iou: {avg_iou}')

当然也可以有节约内存的写法，不必保存所有的(pred, gt)对或者他们的面积信息，在迭代时即完成规约，与调用mean_iou结果相同；

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
from mmseg.core.evaluation.metrics import intersect_and_union, total_area_to_metrics
import numpy as np
from functools import partial
import torch


def memory_efficient():
    num_classes=5
    gen_mat = partial(random_int_matrix, (10, 10), 5)
    cal_areas = partial(intersect_and_union, num_classes=num_classes, ignore_index=-1)

    N = 100
    ps = [gen_mat() for i in range(N)]
    ls = [gen_mat() for i in range(N)]

    total_ai = torch.zeros((num_classes, ), dtype=torch.float64)
    total_au = torch.zeros((num_classes, ), dtype=torch.float64)
    total_ap = torch.zeros((num_classes, ), dtype=torch.float64)
    total_al = torch.zeros((num_classes, ), dtype=torch.float64)
    for p, l in zip(ps, ls):
        ai, au, ap, al = cal_areas(p, l)
        total_ai += ai
        total_au += au
        total_ap += ap
        total_al += al
    results = total_area_to_metrics(total_ai, total_au, total_ap, total_al, metrics=['mIoU'])

    print(results)