python - 使用 numpy 进行高效的加权矢量距离计算

Question

我想计算两组点之间的平方欧氏距离，inputs 和 testing。 inputs 通常是大小为 ~(200, N) 的真实数组，而 testing 通常为 ~(1e8, N)，N 约为 10。距离应为在 N 的每个维度上缩放，所以我将聚合表达式 scale[j]*(inputs[i,j] - testing[ii,j])**2 (其中 scale 是 N 倍的缩放向量)。我正在努力使它尽可能快，特别是因为 N 可能很大。我的第一个测试是

def old_version (inputs, testing, x0):
    nn, d1 = testing.shape
    n, d1 = inputs.shape
    b = np.zeros((n, nn))
    for d in xrange(d1):
        b += x0[d] * (((np.tile(inputs[:, d], (nn, 1)) -
             np.tile (testing[:, d], (n, 1)).T))**2).T
return b

没什么特别的。然后我尝试使用 scipy.spatial.distance.cdist，尽管我仍然需要遍历它以获得正确的缩放

def new_version (inputs, testing, x0):
    # import scipy.spatial.distance as dist
    nn, d1 = testing.shape
    n, d1 = inputs.shape
    b = np.zeros ((n, nn))

    for d in xrange(d1):
        b += x0[d] * dist.cdist(inputs[:, d][:, None], 
             testing[:, d][:, None], 'sqeuclidean')
    return b

看起来 new_version 的缩放比例更好(如 N > 1000)，但我不确定我在这里是否已经尽可能快了。非常感谢任何进一步的想法!

Answer 1

这段代码比你的实现好 10 倍，试一试:

x = np.random.randn(200, 10)
y = np.random.randn(1e5, 10)
scale = np.abs(np.random.randn(1, 10))
scale_sqrt = np.sqrt(scale)

dist_map = dist.cdist(x*scale_sqrt, y*scale_sqrt, 'sqeuclidean')

这些是测试结果:

In [135]: %timeit suggested_version(inputs, testing, x0)

1 loops, best of 3: 341 ms per loop

In [136]: %timeit op_version(inputs, testing, x00) (NOTICE: x00 is a reshape of x0)

1 loops, best of 3: 3.37 s per loop

只要确保在使用较大的 N 时不会出现内存不足的情况。它确实可以减慢速度。