Computing Euclidean distance on 144 dimensions

Computing Euclidean distance on 144 dimensions

去年年底我在博客上看到一篇关于我们的CSAM图像扫描工具.我记得当时在想:这太酷了!图像处理总是很困难的,在Cloudflare部署一个真正的图像识别系统是一个不小的成就!

一段时间后,我在和Kornel聊天:“我们已经完成了图像处理流程中的所有部分,但我们正在努力解决一个组件的性能问题。”扩展到Cloudflare的需求并不容易!

问题在于匹配算法本身的速度。让我详细说明一下。作为约翰在他的博客中解释道,图像匹配算法从处理后的图像中创建模糊哈希。哈希正好是144字节长。例如,可能如下所示:

00e308346a494a188e1043333147267a 653a16b94c33417c12b433095c318012 5612442030d14a4ce82c623f4e224733 1dd84436734e4a5d6e25332e507a8218 6e3b89174e30372d

哈希被设计用于一种模糊匹配算法,可以找到“附近”的相关图像。具体的算法定义得很好,但是要使它更快就要留给程序员了——在Cloudflare,我们需要更快地完成匹配。我们希望将通过我们网络的图像与数百万个已知图像的数据库进行每秒数千次哈希匹配。为了实现这一点,我们需要对匹配算法进行认真的优化。

朴素二次算法

首先想到的算法是O(K*N)复杂性:对于每个查询,遍历数据库中的每个哈希。在这个实现中,创造了很多天真的工作。但到底要做多少工作?

首先,我们需要解释模糊匹配是如何工作的。

给定一个查询哈希,模糊匹配就是数据库中“最近”的哈希。这需要我们定义一个距离。我们将每个散列作为一个包含144个数字的向量,标识144维空间中的一个点。给定两个这样的点,我们可以用标准欧几里德公式计算距离。

但是,对于我们的特定问题,只有当距离小于某个预定义的阈值时,我们才对数据库中的“最近”匹配感兴趣。否则,当距离较大时,我们可以假设图像不相似。这是预期的结果-我们的大多数查询在数据库中没有相关的图像。

这个欧几里德距离算法使用的公式是标准的:

Computing Euclidean distance on 144 dimensions

要计算两个144字节散列之间的距离,我们取每个字节,计算增量,平方,求和为累加器,求平方根,然后打个大!我们有距离!

下面是如何计算C中的平方距离:

Computing Euclidean distance on 144 dimensions

此函数返回平方距离。我们避免计算实际距离以避免运行平方根函数-这很慢。在代码内部,为了性能和简单性,我们主要对平方值进行操作。我们不需要实际的距离值,我们只需要找到最小的向量。在我们的例子中,比较距离还是平方距离无关紧要!

如您所见,模糊匹配基本上是在多维空间中寻找最近点的标准问题。当然,这个问题在过去已经解决了,但我们不要再往前走了。

虽然这段代码可能很简单,但我们希望它相当慢。在一个拥有1M条目的数据库中找到最小的哈希距离,需要遍历所有记录,并且至少需要:

  1. 144*1M减法
  2. 144*1M乘法
  3. 增加144*1M

还有更多。仅此一项,就有4.32亿次手术!在实践中看起来如何?来说明这篇博文我们准备了一套完整的测试套件.已知哈希的大型数据库可以随机数据模拟良好.查询哈希不能是随机的,必须稍微复杂一点,否则这个练习就没有那么有趣了。我们通过对数据库中的实际数据进行字节交换,巧妙地生成了测试——这使我们能够精确地控制测试哈希和数据库哈希之间的距离。查看脚本了解详细信息.这是我们第一次第一个天真的算法:

$make naivey <test-vector.txt./mmdist naive>test vector.tmp 总计:85261.833ms,1536项,每个查询平均55.509ms,18.015 qps

我们在85秒内将1536个测试哈希值与100万个随机向量的数据库进行匹配。找到最近的邻居平均花费了55ms的CPU时间。这对我们的需要来说是相当缓慢的。

SIMD寻求帮助

一个明显的改进是使用更复杂的SIMD指令。SIMD是一种使用一条指令指令CPU处理多个数据点的方法。在处理向量问题时,这是一个完美的策略,我们的任务就是这样。

我们决定使用AVX2,256位向量。我们这样做的原因很简单-更新的AVX版本不支持我们的AMD CPU。另外,在过去,我们对AVX-512的频率缩放不感兴趣.

使用AVX2说起来容易做起来难。没有一条指令可以计算两个uint8向量之间的欧几里德距离!计算两个144字节向量全距离的最快方法用AVX2我们可以找到通过弗拉德:

Computing Euclidean distance on 144 dimensions

它实际上比看上去简单:加载16个字节,将向量从uint8转换为int16,减去向量,将中间和存储为int32,然后重复。最后,我们需要做复杂的4个指令,将部分和提取成最终和。这个AVX2代码性能提高约3倍:

$make naivey-avx2 总计:25911.126ms,1536项,每个查询平均16.869ms,59.280 qps

我们测量了每个项目17ms,仍然低于我们的预期。不幸的是,如果没有重大的改变,我们就无法进一步推进。问题是这段代码受到内存带宽的限制。测量数据来自我的Intel i7-5557U CPU,它的最大理论内存带宽仅为25GB/s。100万个条目的数据库需要137MB,因此至少需要5毫秒才能将数据库输入CPU。有了这个天真的算法,我们就不能低于这个标准了。

优势点树算法

由于天真的暴力方法失败了,我们尝试使用更复杂的算法。我的同事Kornel Lesinïski实施超级酷优势点算法.经过几次反复修改,我们放弃了。我们的问题对这种算法来说异常困难。

我们观察到“维度的诅咒”.空间划分算法不能很好地解决大维数的问题,在我们的例子中,我们有大量的144个维度。K-D树注定要灭亡。位置敏感散列也注定要失败。这是一个奇怪的情况,在这种情况下,空间是难以想象的巨大,但一切都是紧密相连的。空间的体积是一个347位数的长数字,但是点之间的最大距离只有3060平方米(255*255*144)。

空间分割算法是快速的,因为当他们接近找到最近的点时,他们逐渐缩小搜索空间。但是在我们的例子中,公共查询永远不会接近集合中的任何一个点,所以搜索空间不能缩小到有意义的程度。

VP树是一个很有前途的候选树,因为它只对距离进行操作,像二叉树一样将空间划分为近区和远区。当它有一个接近匹配,它可以非常快,不需要访问超过O(对数(N))节点。对于非比赛,它的速度会急剧下降。该算法最终访问了树中几乎一半的节点。在144个维度中,一切都是紧密相连的!尽管该算法避免了访问树中超过一半的节点,但访问剩余节点的成本较高,因此搜索结果总体上较慢。

更聪明的蛮力?

我们在想这个经历。由于空间划分算法不能缩小搜索范围,而且仍然需要检查大量的项目,所以我们应该集中精力快速地检查所有的哈希值。不过,我们必须对内存带宽更加明智——这是以前天真的暴力方法的限制因素。

也许我们不需要从内存中获取所有数据。

短距离

这个突破来自于我们不需要计算散列之间的完整距离的认识。相反,我们只能计算维度的一个子集,比如144个维度中的32个。如果这个距离已经很大,那么就不需要计算完整的距离了!计算更多的点并不会缩短欧几里德距离。

该算法的工作原理如下:

1获取查询哈希并从中提取一个32字节的短哈希

2检查数据库中的100万个32字节的短哈希值。它们必须密密麻麻地堆积在内存中,以允许CPU执行良好的预取,并避免读取我们不需要的数据。

三。如果32字节的短哈希的距离大于或等于目前为止的最佳分数,请继续

4否则,彻底研究散列并计算完整距离。

尽管这个算法需要较少的算术和内存工作,但它并不比先前的简单算法快。看到了吗短接avx2.问题是:我们仍然需要为有希望的散列计算一个完整的距离,而且有相当多的散列。计算有希望的哈希的完整距离会增加足够的工作,包括ALU和内存延迟,以抵消该算法的增益。

有一个细节,我们的特殊应用的图像匹配问题,将有助于我们向前迈进。如前所述,问题不在于找到最近的邻居,而在于证明距离合理的邻居不存在。记住-在实践中,我们不希望找到很多匹配项!我们期望我们输入到算法中的几乎每个图像都与存储在数据库中的图像哈希无关。

我们的算法足以证明在预定的距离阈值内没有邻居存在。假设我们对比220更遥远的散列不感兴趣,220的平方是48400。这使得我们的短距离算法发生了变化工作得更好:

$make short-avx2-threshold 总计:4994.435ms,1536项,每个查询平均3.252ms,307.542qps

原点距离变化

Computing Euclidean distance on 144 dimensions

在某种程度上,John注意到阈值允许额外的优化。我们可以根据散列到某个原点的距离来排序。给定一个起始距离为a的查询散列,我们只能检查距离原点| a-threshold |和| a+threshold |之间的散列。这几乎是每个层次的优势点树的工作原理,只是简单化了。这种优化(按与原点的距离排序数据库中的项目)相对简单,可以帮助我们节省一些工作。

虽然在理论上很好,但这种方法在实践中并没有引入太多的增益,因为向量并不是分组在一起的-它们几乎是随机的!对于我们感兴趣的阈值,原点距离算法的变化给了我们大约20%的速度提升,这是可以的,但不是惊人的。如果我们决定降低阈值,这个更改可能会带来更多的好处,因此对于生产实现来说,这可能是值得的。但是,它不能很好地与查询批处理一起工作。

所以不是:

  1. [a1、a2、a3],
  2. [b1、b2、b3],
  3. [c1、c2、c3],

...

我们可以把它放在记忆里转置:

  1. [a1、b1、c1],
  2. [a2、b2、c2],
  3. [a3、b3、c3],

...

现在我们可以使用一个内存操作加载16个第一个字节的散列。在下一步中,我们可以使用一条指令减去查询哈希的第一个字节,依此类推。算法与上面定义的完全相同;我们只是使AVX的数据更易于加载和处理。

热循环代码甚至看起来相对漂亮:

Computing Euclidean distance on 144 dimensions

通过调整好的批量大小和短距离尺寸参数,我们可以查看此算法的性能:

$make short-inv-avx2 总计:1118.669ms,1536项,每个查询平均0.728ms,1373.062 qps

哇!这真是太棒了。我们已经从55ms开始优化每页的内存,或者说减少了73毫秒的内存。

Computing Euclidean distance on 144 dimensions

丹尼斯·巴赫瓦洛夫书中的屋顶模型‌‌

如果您对这样的体系结构调优感兴趣,请看一下丹尼斯·巴赫瓦洛夫的新表演书.它讨论了屋顶线模型分析,这和我们在这里做的差不多。

一定要看看我们的代码告诉我们,如果我们错过了一些优化!

摘要

多棒的优化之旅!我们在内存和ALU瓶颈代码之间跳跃。我们讨论了更复杂的算法,但最终,一个蛮力算法-虽然调整-给了我们最好的结果。

为了得到更好的数字,我用CUDA对Nvidia GPU进行了实验。CUDA的本质是vabsdiff4型dp4a型完美地解决了这个问题。

V100给了我们一些惊人的数字,但我并不完全满意。考虑到有多少AMD Ryzen核与AVX2我们可以得到的成本,一个服务器级的GPU,我们倾向于通用计算这一特殊问题。

这是我们每天处理的复杂问题的一个很好的例子。即使是最好的技术在“Cloudflare规模”下工作也需要跳出框框思考。有时我们在找到最优解之前,会重写几十次,有时我们会采用蛮力算法,非常优化。

需要大量CPU运算和哈希运算的图像匹配问题非常具有挑战性。我们在边缘上可用的CPU非常稀少,这样的工作负载非常昂贵。即使在这篇博客文章中谈到了优化工作,大规模运行CSAM扫描仪也是一个挑战,需要大量的工程工作。我们还没说完!在我们满意之前,我们需要解决更多的难题。