python - 加速某些 numpy 操作的有效方法

Question

我正在尝试找到一个高效的代码，而不是下面的代码(这只是我的代码的一部分)，以提高速度:

for pr in some_list:
    Tp = T[partition[pr]].sum(0)
    Tpx = np.dot(Tp, xhat)
    hp = h[partition[[pr]].sum(0)
    up = (uk[partition[pr][:]].sum(0))/len(partition[pr])
    hpu = hpu + np.dot(hp.T, up)
    Tpu = Tpu + np.dot(Tp.T, up)

我至少还有两个类似的代码块。正如你所看到的，我使用了三次花式索引(真的找不到其他方法)。在我的算法中，我需要非常快地完成这部分，但现在还没有发生。我将非常感谢任何建议。

谢谢大家。

最好，

Answer 1

如果您的分区很少并且每个分区都有很多元素，您应该考虑交换对象的索引。对形状数组 (30,1000) 求和沿着它的第二个维度应该比对形状数组 (1000,30) 求和更快沿着它的第一个维度，因为在前一种情况下，您总是为每个剩余索引求和连续的内存块(即 arr[k,:] 对于每个 k )。因此，如果您将求和索引放在最后(并在执行时删除一些尾随的单例维度)，您可能会获得加速。

如hpaulj noted in a comment ，目前尚不清楚如何对循环进行矢量化。但是，由于它对性能至关重要，您仍然可以尝试对部分工作进行矢量化。

我建议您存储hp , up和Tp对于每个分区(在预分配之后)，然后在单个向量化步骤中执行标量/矩阵乘积。另请注意 Tpx在您的示例中未使用，因此我在这里省略了它(无论您用它做什么，您都可以像其他示例一样进行操作):

part_len = len(some_list) # number of partitions, N
Tpshape = (part_len,) + T.shape[1:] # (N,30,100) if T was (1000,30,100)
hpshape = (part_len,) + h.shape[1:] # (N,30,1) if h was (1000,30,1)
upshape = (part_len,) + uk.shape[1:] # (N,30,1) if uk was (1000,30,1)
Tp = np.zeros(Tpshape)
hp = np.zeros(hpshape)
up = np.zeros(upshape)

for ipr,pr in enumerate(some_list):
    Tp[ipr,:,:] = T[partition[pr]].sum(0)
    hp[ipr,:,:] = h[partition[[pr]].sum(0)
    up[ipr,:,:] = uk[partition[pr]].sum(0)/len(partition[pr])

# compute vectorized dot products:
#Tpx unclear in original, omitted
# sum over second index (dot), sum over first index (sum in loop)
hpu = np.einsum('abc,abd->cd',hp,up)  # shape (1,1)
Tpu = np.einsum('abc,abd->cd',Tp,up)  # shape (100,1)

显然关键人物是 numpy.einsum 。当然如果 hpu和Tpu在循环之前有一些先前的值，您必须使用 einsum 的结果来增加这些值如上所述。

至于einsum ，它对任意维度的数组执行求和和收缩。上面出现的模式，'abc,abd->cd' ，当应用于 3d 数组时 A和B ，将返回一个二维数组 C ，具有以下定义(数学伪代码):

C(c,d) = sum_a sum_b A(a,b,c)*B(a,b,d)

对于给定的修复a求和索引，里面的内容是

sum_b A(a,b,c)*B(a,b,d)

其中，如果 c和d索引被保留，将相当于 np.dot(A(a,:,:).T,B(a,:,:)) 。因为我们根据 a 对这些矩阵求和同样，我们应该完全按照你的循环版本所做的事情，将每个np.dot()相加。总金额的贡献。