如何提高Python(NumPy)中两个大型批处理数组的交叉二元乘积的性能

Numpy的问题在于，默认情况下会创建许多相对较大的临时数组(为了方便起见).这太贵了.你可以做in-place operations个，尽管这做起来很麻烦.此外，您还可以对最终的乘法进行因式分解，以确定条件是否成立:

下面是一个例子:

def cdya_opt(A, B):
    i, j = [a.ravel() for a in np.indices((6, 6))]

    a = np.array([(0, 0), (1, 1), (2, 2), (0, 1), (1, 2), (0, 2)])
    b = np.array([0, 3, 5, 3, 1, 4, 5, 4, 2]).reshape(3, 3)

    i, j, k, l = np.hstack([a[i], a[j]]).T

    ik = b[i, k].reshape(6, 6)
    jl = b[j, l].reshape(6, 6)

    il = b[i, l].reshape(6, 6)
    kj = b[k, j].reshape(6, 6)

    x1, x2, x3, x4 = A[ik], B[jl], A[il], B[kj]
    np.multiply(x1, x2, out=x1)
    np.multiply(x3, x4, out=x3)
    np.add(x1, x3, out=x1)

    if A is not B:
        x5, x6, x7, x8 = B[ik], A[jl], B[il], A[kj]
        np.multiply(x5, x6, out=x5)
        np.multiply(x7, x8, out=x7)
        np.add(x5, x7, out=x5)
        np.add(x5, x1, out=x1)
        np.multiply(x1, 0.25, out=x1)
    else:
        np.multiply(x1, 0.5, out=x1)

    return x1

如果您可以使用100之类的其他包，那么您可以编写更快的实现:

import numexpr as ne

def cdya_opt2(A, B):
    i, j = [a.ravel() for a in np.indices((6, 6))]

    a = np.array([(0, 0), (1, 1), (2, 2), (0, 1), (1, 2), (0, 2)])
    b = np.array([0, 3, 5, 3, 1, 4, 5, 4, 2]).reshape(3, 3)

    i, j, k, l = np.hstack([a[i], a[j]]).T

    ik = b[i, k].reshape(6, 6)
    jl = b[j, l].reshape(6, 6)

    il = b[i, l].reshape(6, 6)
    kj = b[k, j].reshape(6, 6)

    x1, x2, x3, x4 = A[ik], B[jl], A[il], B[kj]

    if A is not B:
        x5, x6, x7, x8 = B[ik], A[jl], B[il], A[kj]
        ne.evaluate('(x1 * x2 + x3 * x4 + x5 * x6 + x7 * x8) * 0.25', out=x1)
    else:
        ne.evaluate('(x1 * x2 + x3 * x4) * 0.5', out=x1)

    return x1

性能结果

以下是在我的机器(i5-9600KF CPU，运行在Windows上，安装了Numpy 1.24.3)上的结果:

cdya:       112.0 ms
cdya_opt:    70.7 ms
cdya_opt2:   59.1 ms  <-----

上一个版本的速度几乎快了一倍.大部分时间都花在像A[ik]这样的操作上，而这些操作在Numpy中并没有得到很好的优化.使用Cython(或Numba)是使这段代码更快的好方法.请注意，与Numba相反，Cython不是JIT.