使用numba jit、Python的段间距离

在过去的一周里，我一直在问这个堆栈上的相关问题，试图找出在Python中使用@jit decorator和Numba时我不理解的地方。然而，我遇到了麻烦，所以我只写下整个问题

当前的问题是计算大量段对之间的最小距离。线段在三维中由其起点和终点表示。从数学上讲，每个线段都被参数化为[AB]=A+（B-A）*s，其中[0,1]中的s，A和B是线段的起点和终点。对于两个这样的段，可以计算最小距离并给出公式

我已经在另一个网站上暴露了这个问题，给出的答案是通过矢量化问题来替换代码的双循环，但是这会影响到大型段集的内存问题。因此，我决定坚持使用循环，并改用numba的jit

由于这个问题的解决方案需要大量的点产品，而numpy的点产品是，我从实现自己的3D点产品开始

import numpy as np
from numba import jit, autojit, double, float64, float32, void, int32

def my_dot(a,b):
    res = a[0]*b[0]+a[1]*b[1]+a[2]*b[2]
    return res

dot_jit = jit(double(double[:], double[:]))(my_dot)    #I know, it's not of much use here.

计算N段中所有对的最小距离的函数将Nx6数组（6个坐标）作为输入

因此，compute_stuff（arg）将一个2D np.array（double[：，：]）作为参数，执行一系列numba支持的（？）操作，并返回另一个2D np.array（double[：，：]）。但是,

v = np.random.random( (100,6) )
%timeit compute_stuff(v)
%timeit fast_compute_stuff(v)

---

我得到134和123毫秒每循环。你能解释一下为什么我不能加快我的功能吗？如果您有任何反馈，我们将不胜感激。

以下是我的代码版本，它的速度要快得多：

@jit(nopython=True)
def dot(a,b):
    res = a[0]*b[0]+a[1]*b[1]+a[2]*b[2]
    return res

@jit
def compute_stuff2(array_to_compute):
    N = array_to_compute.shape[0]
    con_mat = np.zeros((N,N))

    p0 = np.zeros(3)
    p1 = np.zeros(3)
    q0 = np.zeros(3)
    q1 = np.zeros(3)

    p0m1 = np.zeros(3)
    p1m0 = np.zeros(3)
    q0m1 = np.zeros(3)
    q1m0 = np.zeros(3)
    p0mq0 = np.zeros(3)

    for i in range(N):
        for j in range(i+1,N):

            for k in xrange(3):
                p0[k] = array_to_compute[i,k]
                p1[k] = array_to_compute[i,k+3]
                q0[k] = array_to_compute[j,k]
                q1[k] = array_to_compute[j,k+3]

                p0m1[k] = p0[k] - p1[k]
                p1m0[k] = -p0m1[k]

                q0m1[k] = q0[k] - q1[k]
                q1m0[k] = -q0m1[k]

                p0mq0[k] = p0[k] - q0[k]

            s = ( dot(p1m0, q1m0)*dot(q1m0, p0mq0) - dot(q1m0, q1m0)*dot(p1m0, p0mq0))/( dot(p1m0, p1m0)*dot(q1m0, q1m0) - dot(p1m0, q1m0)**2 )
            t = ( dot(p1m0, p1m0)*dot(q1m0, p0mq0) - dot(p1m0, q1m0)*dot(p1m0, p0mq0))/( dot(p1m0, p1m0)*dot(q1m0, q1m0) - dot(p1m0, q1m0)**2 )


            for k in xrange(3):
                con_mat[i,j] += (p0[k]+(p1[k]-p0[k])*s-(q0[k]+(q1[k]-q0[k])*t))**2 

    return con_mat

以及时间安排：

In [38]:

v = np.random.random( (100,6) )
%timeit compute_stuff(v)
%timeit fast_compute_stuff(v)
%timeit compute_stuff2(v)

np.allclose(compute_stuff2(v), compute_stuff(v))

#10 loops, best of 3: 107 ms per loop
#10 loops, best of 3: 108 ms per loop
#10000 loops, best of 3: 114 µs per loop
#True

我加快这一进程的基本策略是：

• 去掉所有数组表达式并显式展开矢量化（尤其是因为数组太小）
• 在调用

  dot

  方法时，消除冗余计算（减去两个向量）
• 将所有数组创建移到嵌套for循环的外部，以便numba可以执行一些操作。这还避免了创建许多成本高昂的小型阵列。最好分配一次并重用内存

另一件需要注意的事情是，对于numba的最新版本，过去被称为

autojit

（即，让numba对输入进行类型推断）并且现在只是没有类型提示的装饰器，通常与根据我的经验指定输入类型一样快

---

此外，在OS X上使用Nanconda python发行版和python 2.7.9，使用numba 0.17.0运行计时。

使用numba的JIT编译器，您不太可能击败

np.dot

np.dot

只是一个调用BLAS

*gemm/*gemv

函数的薄包装器，这些函数经过了大量优化，通常是多线程的。您最好的选择可能是确保numpy链接到您可以使用的最快的多线程BLAS库（可能是Intel的MKL或OpenBLAS）。问题不是打败np.dot，而是如果jit编译器遇到np.dot调用，它不能推断它的返回类型，也不会加快我的整个函数的速度（顺便说一句，我编写的dot_jit比np.dot对于三维矢量标量产品要快），您对原始代码进行了线条分析了吗？我怀疑您的大部分时间都花在

np.dot

中，因此不应该期望从嵌套的

for

循环中获得多少性能好处。使用cProfile，我发现对于1000个段，我在这些深度操作（np.dot、np.sum等）中的累积时间大约是13秒中的1秒，再看一遍你的代码，我意识到这是因为你点的向量只有2长！你能在你的问题中公布测线时间吗？我刚刚测试了5000段的数组。它将执行时间从6分钟减少到343毫秒。我读了一些关于你提供的环提升的链接，我想除了你策略的第一点之外，我什么都懂。为什么作为循环标量操作显式执行数组操作更快？由于您似乎对这个主题非常熟悉，对于更大的阵列，您认为将其实现为cuda（numbapro）会进一步加快速度吗？我没有通过numbapro将numba与cuda一起使用，因此我无法对此发表评论。就我的策略的第一点而言，目前numba在将UFUNC编译为本机代码方面有一定的限制，这基本上涉及传入输出数组：。因此，您的代码可能可以实现这一点，但我选择手动展开矢量化。

In [38]:

v = np.random.random( (100,6) )
%timeit compute_stuff(v)
%timeit fast_compute_stuff(v)
%timeit compute_stuff2(v)

np.allclose(compute_stuff2(v), compute_stuff(v))

#10 loops, best of 3: 107 ms per loop
#10 loops, best of 3: 108 ms per loop
#10000 loops, best of 3: 114 µs per loop
#True