python - numba 输出的差异

Question

我在学习工作中实现了基本的最近邻搜索。事实上，基本的 numpy 实现运行良好，但只是添加“@jit”装饰器(在 Numba 中编译)，输出是不同的(由于某些未知原因，它最终复制了一些邻居......)

基本算法如下:

import numpy as np
from numba import jit

@jit(nopython=True)
def knn(p, points, k):
    '''Find the k nearest neighbors (brute force) of the point p
    in the list points (each row is a point)'''

    n = p.size  # Lenght of the points
    M = points.shape[0]  # Number of points
    neighbors = np.zeros((k,n))
    distances = 1e6*np.ones(k)

    for i in xrange(M):
        d = 0
        pt = points[i, :]  # Point to compare
        for r in xrange(n):  # For each coordinate
            aux = p[r] - pt[r]
            d += aux * aux
        if d < distances[k-1]:  # We find a new neighbor
            pos = k-1
            while pos>0 and d<distances[pos-1]:  # Find the position
                pos -= 1
            pt = points[i, :]
            # Insert neighbor and distance:
            neighbors[pos+1:, :] = neighbors[pos:-1, :]
            neighbors[pos, :] = pt
            distances[pos+1:] = distances[pos:-1]
            distances[pos] = d

    return neighbors, distances

用于测试:

p = np.random.rand(10)
points = np.random.rand(250, 10)
k = 5
neighbors = knn(p, points, k)

没有@jit 装饰器，得到正确答案:

In [1]: distances
Out[1]: array([ 0.3933974 ,  0.44754336,  0.54548715,  0.55619749,  0.5657846 ])

但是 Numba 编译给出了奇怪的输出:

Out[2]: distances
Out[2]: array([ 0.3933974 ,  0.44754336,  0.54548715,  0.54548715,  0.54548715])

有人可以帮忙吗？我不明白为什么会这样......

谢谢。

Answer 1

我认为问题是 Numba 处理将一个切片写入另一个切片时，当这些切片重叠时与不重叠时运行时的处理方式不同。我不熟悉 numpy 的内部结构，但也许有特殊的逻辑来处理像这样的 volatile 内存操作，而 Numba 中没有。更改以下行，jit 装饰器的结果与普通 python 版本一致:

neighbors[pos+1:, :] = neighbors[pos:-1, :].copy()
...
distances[pos+1:] = distances[pos:-1].copy()