提高距离计算速度的建议

Question

考虑一下下面这个类：

class SquareErrorDistance(object):
    def __init__(self, dataSample):
        variance = var(list(dataSample))
        if variance == 0:
            self._norm = 1.0
        else:
            self._norm = 1.0 / (2 * variance)

    def __call__(self, u, v): # u and v are floats
        return (u - v) ** 2 * self._norm

我用它来计算向量中两个元素之间的距离。基本上，我为使用这种距离测量的向量的每个维度创建一个这个类的实例（有些维度使用其他的距离测量）。性能分析显示，这个类的__call__函数占了我knn实现运行时间的90%（谁能想到呢）。我觉得用纯Python的方式来加速这个过程是不太可能的，但如果我用C来实现呢？

如果我运行一个简单的C程序，仅仅使用上面的公式计算随机值的距离，它的速度比Python快得多。所以我尝试使用ctypes来调用一个进行计算的C函数，但显然参数和返回值的转换成本太高，导致生成的代码反而更慢。

当然，我可以把整个knn实现用C写出来，然后直接调用，但问题是，正如我所说，我对向量的某些维度使用不同的距离函数，把这些转换成C代码会太麻烦。

那么我还有什么其他选择呢？使用Python C-API编写C函数能消除这些额外开销吗？有没有其他方法可以加速这个计算？

Answer 1

这可能帮助不大，但你可以试着用嵌套函数来重写它：

def SquareErrorDistance(dataSample):
    variance = var(list(dataSample))
    if variance == 0:
        def f(u, v):
            x = u - v
            return x * x
    else:
        norm = 1.0 / (2 * variance)
        def f(u, v):
            x = u - v
            return x * x * norm
    return f

Answer 2

下面的这段cython代码（我知道__init__的第一行不一样，我随便换了些东西，因为我不知道var是什么，而且这也无所谓——你提到过__call__是性能瓶颈）：

cdef class SquareErrorDistance:
    cdef double _norm

    def __init__(self, dataSample):
        variance = round(sum(dataSample)/len(dataSample))
        if variance == 0:
            self._norm = 1.0
        else:
            self._norm = 1.0 / (2 * variance)

    def __call__(self, double u, double v): # u and v are floats
        return (u - v) ** 2 * self._norm

通过一个简单的setup.py编译（就像文档里的例子，只是改了文件名），它在一个简单的timeit基准测试中，性能几乎比纯Python快20倍。值得注意的是，唯一的变化是_norm字段和__call__参数的cdef声明。我觉得这非常令人印象深刻。