在 Python 中优化两个矩阵的直方图距离度量

Question

我有两个矩阵A 和B，每个矩阵的大小为NxM，其中N 是样本数，M 是直方图箱的大小。因此，每一行代表该特定样本的直方图。

我想做的是为不同的一对样本计算两个矩阵之间的chi-square 距离。因此，矩阵A 中的每一行将与另一个矩阵B 中的所有行进行比较，从而得到一个大小为NxN 和C[i,j] 的最终矩阵C[i,j] 对应于chi-square 之间的距离A[i] 和 B[j] 直方图。

这是我完成这项工作的python代码：

def chi_square(histA,histB):
   esp = 1.e-10
   d = sum((histA-histB)**2/(histA+histB+eps))
   return 0.5*d
def matrix_cost(A,B):
   a,_ = A.shape
   b,_ = B.shape
   C = zeros((a,b))
   for i in xrange(a):
      for j in xrange(b):
         C[i,j] = chi_square(A[i],B[j])
  return C

目前，对于100x70 矩阵，整个过程需要 0.1 秒。

有什么办法可以提高这种性能吗？

如果有任何想法或建议，我将不胜感激。

谢谢。

Answer 1

当然！我假设您使用的是numpy？

如果您有可用的 RAM，您可以使用 broadcast 数组并使用 numpy 对这些数组的操作进行高效矢量化。

方法如下：

Abroad = A[:,np.newaxis,:]  # prepared for broadcasting
C = np.sum((Abroad - B)**2/(Abroad + B), axis=-1)/2.

与您的算法相比，我的平台上的时间考虑因素显示了 10 倍的速度增益。

比前一个选项使用更少 RAM 的较慢选项（但仍比您的原始算法更快）只是将 A 的行广播到 2D 数组中：

def new_way(A,B):
    C = np.empty((A.shape[0],B.shape[0]))
    for rowind, row in enumerate(A):
        C[rowind,:] = np.sum((row - B)**2/(row + B), axis=-1)/2.
    return C

这样做的好处是它可以针对形状 (N,M) 远大于 (100,70) 的数组运行。

如果您没有可用的内存，您也可以查看 Theano 将昂贵的 for 循环推送到 C 级别。与 (100,70) 数组和 (1000,70) 的第一个选项（不考虑初始编译时间）相比，我获得了 2 倍的速度提升：

import theano
import theano.tensor as T
X = T.matrix("X")
Y = T.matrix("Y")
results, updates = theano.scan(lambda x_i: ((x_i - Y)**2/(x_i+Y)).sum(axis=1)/2., sequences=X)
chi_square_norm = theano.function(inputs=[X, Y], outputs=[results])
chi_square_norm(A,B)  # same result