Python计算numpy数组向量4元素乘积均值的效率优化问题

你现有代码效率低的核心原因有两个：

• 逻辑冗余：np.einsum('ni,nj,nk,nr->ijkr',array,array,array,array)已经自动对样本轴n做了求和，输出是形状为(30,30,30,30)的四阶张量，后续添加的sum(axis=0)完全多余，不仅会得到错误的计算结果，还额外增加了计算开销。
• 未开启张量缩并优化：numpy自带的einsum默认关闭最优计算路径选择，四阶张量缩并的原生实现会产生5e5 * 30^4 = 4.05e11次浮点运算，耗时自然极高。

方案1：修正einsum用法，开启内置优化 #

无需引入额外依赖，仅修改参数即可获得10~100倍的速度提升，适配CPU场景：

import numpy as np

# array.shape = (500000, 30)
expect = np.einsum('ni,nj,nk,nr->ijkr', array, array, array, array, optimize='optimal') / array.shape[0]

optimize='optimal'参数会让numpy自动搜索最优的张量缩并路径，大幅降低实际运算量。

方案2：仅计算无重复4元素组合（最大幅提速） #

如果你需求中的“选取任意4个元素”指的是不重复、不计顺序的组合（而非考虑索引顺序的全排列），30个元素选4个仅27405种组合，相比30^4=81万种排列计算量下降97%，速度提升最明显：

import numpy as np
from itertools import combinations

n_samples, n_dim = array.shape
expect = np.zeros((n_dim, n_dim, n_dim, n_dim))
for i, j, k, r in combinations(range(n_dim), 4):
    prod = array[:, i] * array[:, j] * array[:, k] * array[:, r]
    expect[i, j, k, r] = prod.mean()
    # 若需要张量对称性，可补充其余索引位置的赋值逻辑

方案3：高性能张量库加速 #

如果有GPU资源，可使用优化效果更好的opt_einsum或者JAX进一步提速：

import opt_einsum as oe

expect = oe.contract('ni,nj,nk,nr->ijkr', array, array, array, array) / array.shape[0]

内容的提问来源于stack exchange，提问作者Inm