使用jax加速Hamming Distance的计算

技术背景

一般认为Jax是谷歌为了取代TensorFlow而推出的一款全新的端到端可微的框架，但是Jax同时也集成了绝大部分的numpy函数，这就使得我们可以更加简便的从numpy的计算习惯中切换到GPU的计算中。Jax除了支持GPU的张量运算，更重要的一个方面是Jax还支持谷歌自己的硬件TPU的张量运算。关于张量计算，可以参考前面写过的这一篇博客。

而标题中的另外一个概念：Hamming Distance是用来衡量两个字符串之间的相似关系评分算法，如果两个字符串的所有元素完全相同，那么就会得到一个0的分数，如果两个长度各为100的字符串完全不相同（即每一个位置的字符都完全不同），那么得到的Hamming Distance就是100。而关于Normalized Hamming Distance的概念，则是为了使得结果更加的收敛，因此在Hamming Distance的基础之上再除以字符串的总长度，得到一个新的评分。举个例子说，Boy和Bob这两个字符串的Hamming Distance为1，而Normalized Hamming Distance为\(\frac{1}{3}\)。

Numpy和Jax代码实现

一般计算Hamming Distance可以通过scipy中自带的distance.hamming来计算两个字符串之间的相似度，然而我们在日常的计算中更多的会把字符串转化成一个用数字来表示的数组，因此这里我们可以直接使用numpy的equal函数之后在做一个sum即可得到我们需要的Hamming Distance，如果再除以一个数组长度，那么就是Normalized Hamming Distance。由于Jax上实现了GPU版本的Numpy的函数，因此这里我们将Numpy的函数和Jax的函数写到一起来进行对比，尤其是时间上的一个衡量。这里测试的逻辑是：我们先通过Numpy来生成两个给定维度的随机数，然后将其转化成两个Jax格式的数组，然后分别对这两组不同格式的数组分别用Numpy和Jax计算Hamming Distance，最终统计多次运行所得到的时间。

# normalized_hamming_distance.py
import numpy as np
import jax.numpy as jnp
import time

if __name__ == '__main__':
    np.random.seed(1)
    length = 100000000
    arr1 = np.random.randint(5, size=(length,),dtype=np.int32)
    arr2 = np.random.randint(5, size=(length,),dtype=np.int32)
    arr1_jax = jnp.array(arr1)
    arr2_jax = jnp.array(arr2)
    # Start Testing
    time0 = time.time()
    for _ in range(10):
        nhd = np.sum(np.equal(arr1,arr2))/length
    time1 = time.time()
    for _ in range(10):
        nhd_jax = jnp.sum(jnp.equal(arr1_jax,arr2_jax))/length
    time2 = time.time()
    # Result analysis
    print ('The normalized hamming distance by numpy is: {}'.format(nhd))
    print ('The normalized hamming distance by jax is: {}'.format(nhd_jax))
    print ('The time cost by numpy is: {}s'.format(time1-time0))
    print ('The time cost by jax is: {}s'.format(time2-time1))

输出结果如下所示：

The normalized hamming distance by numpy is: 0.20006858
The normalized hamming distance by jax is: 0.20006857812404633
The time cost by numpy is: 1.7030510902404785s
The time cost by jax is: 0.28351473808288574s

经过对比，我们发现Jax所实现的Numpy的GPU版本，可以在几乎不用改动接口的条件下，极大程度上的加速了Numpy的计算过程。

总结概要

本文通过对比Jax和Numpy计算Normalized Hamming Distance的过程来对比了Jax所实现的Numpy的GPU版本所带来的加速效果。实际上在维度比较小的时候，Numpy还是有非常轻量级的优势，此时GPU的加速效果并没有很好的体现出来。但是在规模较大的输入场景下，GPU的并行加速效果简直无敌，而且几乎没有改动原本Numpy的函数接口。除此之外，Jax作为一个函数式编程的端到端可微编程框架，支持jit、vmap、pmap和xmap等非常神奇的加速和并行化功能，为深度学习等领域提供了非常强有力的支持。