矩阵计算还是循环计算？一次空间换时间的优化

平时开发时，为了开发速度以及可读性，我们常常会直接选择写for循环来实现计算逻辑。但是熟悉科学计算的朋友一定知道，这样的代码在计算效率上并非最优的，在某些时候我们可以仅仅通过将计算逻辑改造为矩阵(张量)计算以实现计算效率的明显提升。本文就以笔者之前做表格还原任务时遇到的一个性能问题为例，对比一下算法服务中矩阵计算与循环计算的效率。

1.性能问题的暴露

由于测试阶段大部分测试样本都是单元格数量较少的小型表格，并未发现任何问题。有一天，某业务线的业务数据中包含了一张稍显密集的表格图片（51行9列），笔者出于好奇拿来测试了一下表格识别接口，结果发现响应速度与测试阶段相比明显变慢了。确认服务器无异常后，笔者对服务的计算流程中的每一步进行了耗时分析。结果如下图所示

嗯？怎么TSR的耗时如此离谱？该服务在对有线表格的处理逻辑中，在TSR任务上使用了modelscope的读光模型(Cycle-CenterNet)，推理方式为modelscope官方pipeline。难道是官方代码有问题吗？于是笔者继续对modelscope的TableRecognitionPipeline中的计算过程进行了耗时分析。

从上图可以看到，该pipeline中耗时最高的两个计算分别是模型推理与后处理（尤其是group_bbox_by_gbox函数）。其中模型推理耗时较高是因为笔者的测试环境非gpu服务器，因此进行的的是cpu推理，这一点在生产环境的gpu推理条件下可以得到彻底解决。由此可以断言，性能问题集中于后处理模块中的group_bbox_by_gbox函数中。

2.问题的定位

现在我们已经知道group_bbox_by_gbox的性能堪忧了，那么我们接下来就要继续探索2个问题：

1. 这个函数究竟在算什么，为何计算得如此缓慢？
2. 有没有计算速度更快的实现方法？

我们先来研究第一个问题——这个函数涉及哪些计算。下面是该函数的代码，源码位于modelscope(1.3.2)中的table_process.py

plaintext

def group_bbox_by_gbox(bboxes,
                       gboxes,
                       score_thred=0.3,
                       v2c_dist_thred=2,
                       c2v_dist_thred=0.5):

    def point_in_box(box, point):
        x1, y1, x2, y2 = box[0], box[1], box[2], box[3]
        x3, y3, x4, y4 = box[4], box[5], box[6], box[7]
        ctx, cty = point[0], point[1]
        a = (x2 - x1) * (cty - y1) - (y2 - y1) * (ctx - x1)
        b = (x3 - x2) * (cty - y2) - (y3 - y2) * (ctx - x2)
        c = (x4 - x3) * (cty - y3) - (y4 - y3) * (ctx - x3)
        d = (x1 - x4) * (cty - y4) - (y1 - y4) * (ctx - x4)
        if (a > 0 and b > 0 and c > 0 and d > 0) or (a < 0 and b < 0 and c < 0
                                                     and d < 0):
            return True
        else:
            return False

    def get_distance(pt1, pt2):
        return math.sqrt((pt1[0] - pt2[0]) * (pt1[0] - pt2[0])
                         + (pt1[1] - pt2[1]) * (pt1[1] - pt2[1]))

    dets = copy.deepcopy(bboxes)
    sign = np.zeros((len(dets), 4))

    for idx, gbox in enumerate(gboxes):  # vertex x,y, gbox, score
        if gbox[10] < score_thred:
            break
        vertex = [gbox[0], gbox[1]]
        for i in range(0, 4):
            center = [gbox[2 * i + 2], gbox[2 * i + 3]]
            if get_distance(vertex, center) < v2c_dist_thred:
                continue
            for k, bbox in enumerate(dets):
                if bbox[8] < score_thred:
                    break
                if sum(sign[k]) == 4:
                    continue
                w = (abs(bbox[6] - bbox[0]) + abs(bbox[4] - bbox[2])) / 2
                h = (abs(bbox[3] - bbox[1]) + abs(bbox[5] - bbox[7])) / 2
                m = max(w, h)
                if point_in_box(bbox, center):
                    min_dist, min_id = 1e4, -1
                    for j in range(0, 4):
                        dist = get_distance(vertex,
                                            [bbox[2 * j], bbox[2 * j + 1]])
                        if dist < min_dist:
                            min_dist = dist
                            min_id = j
                    if (min_id > -1 and min_dist < c2v_dist_thred * m
                            and sign[k][min_id] == 0):
                        bboxes[k][2 * min_id] = vertex[0]
                        bboxes[k][2 * min_id + 1] = vertex[1]
                        sign[k][min_id] = 1
    return bboxes

不得不承认这段代码的可读性确实很强，除函数逻辑外，我们还可以大致猜测到gbox、bbox的向量结构如下

而函数逻辑中，主要的计算集中在4层嵌套的for循环里。主要涉及4个计算：

1. 阈值判断（大小比较）
2. 根据bbox四个顶点的坐标计算bbox的宽(w)和高(h)
3. 计算vertex与center的距离
4. 计算vertex与bbox四个顶点的距离

看到这里，我们会发现这里的计算都可以归结为矩阵线性运算、数组切片、条件赋值，而这些计算几乎都能在矩阵/数组/张量层面进行。这一发现预示着这个4层for循环也许可以通过矩阵计算的形式来进行优化。

3.优化方案

惭愧地说，笔者到最后也没能使用矩阵计算去完全替代这一个4层for循环。原因在于没能很好地用矩阵运算来描述其中的point_in_box和部分if-else判断逻辑。但回顾上述4个计算，2、3、4三个计算可以通过预先算好对应数值，在for循环中直接取值进行逻辑判断的方案来进行优化。

确定好优化的思路以后，我们就可以开始动手优化了。然而优化的过程，也不能说是一帆风顺的。

以vertex与center的距离计算为例，假设vertex个数为$n$ 、center个数为$m$，那么我们输入的是一个$n \times 2$的矩阵$V$与一个$m \times 2$的矩阵$C$，期望输出一个$n \times m$的矩阵$D$，其中$D_{ij}$表示$vertex_i$与$center_j$的欧氏距离。

二维向量的欧氏距离计算过程非常简单

$$d(\textbf{x},\textbf{y}) = \sqrt{\sum_{i=1}^{2}(x_{i}-y_{i})}=\sqrt{(x_0-y_0)^2 + (x_1-y_1)^2}$$

但稍加观察就会发现，我们并不能直接使用$V$与$C$计算得到$D$，难点在于公式中的减法运算。我们知道，矩阵乘法的本质是$\sum\limits_{i=1}^{n}{x_i \cdot y_i}$，但对于欧氏距离中的“两个相同维度的向量，相同位置相减”的计算，却显然不能使用这种方法来计算；而矩阵加减法，仅限于shape相同的两个矩阵，我们这里无法保证$V$与$C$的shape相同，因此也无法直接使用矩阵加减运算。

这里，笔者用了一种比较笨的办法来处理——既然矩阵加减法的限制在于shape不同，那么我们只要能够保证做减法的两个矩阵的shape一致不就可以了吗？上文我们说过，目标输出是一个$n \times m$的矩阵，若做减法的两个矩阵的shape是$nm \times 2$，则可顺利实现欧氏距离的矩阵计算，并且后续可通过reshape来恢复$n \times m$的shape以满足输出要求。要实现这个操作，只需要对$V$和$C$以不同方式repeat即可。

4.两种方案的性能对比

值得注意的一点是，上述方案并不能减少计算量，而是希望利用numpy、pytorch、tensorflow等科学计算框架对于矩阵计算的针对性优化来提升该函数的计算速度。原始方案虽然被python的for循环效率所拖累，但作了若干个剪枝判断；而上述方案由于包含repeat操作，因此引入了额外的计算量。那么，该方案是否有效呢？下面两张图给出了答案

可以看到，优化后，group_bbox_by_gbox函数的耗时由原先的4.58秒缩短至2.71秒，速度提升了40%左右。且根据笔者的多次验证，该提升幅度基本稳定。

后记

工程优化无非就是空间与时间的优化，根据工程的目的不同、优化方案的选择自然也不同，适合的即是最好的。该方案只是工程部署中的一次针对性优化方案，其劣势也非常明显：

1. 牺牲了代码的可读性（虽然这很大概率是由于笔者个人代码水平导致的）；
2. 不必要的空间浪费。因为我们需要事先创建2个$n \times m \times 2$的张量，但张量中的每个值在下游计算中可能只会用到一次，这无疑是空间资源的一种浪费。

当然，该方案也还存在很多可继续优化的点，例如张量计算能使用gpu加速等等，笔者也会持续去尝试更多更好的工程实现方案。