通过3-d矩阵的每个切片将2-d矩阵的每列相乘的更有效方式

时间:2017-07-07 21:59:19

标签: python performance numpy matrix linear-algebra

我有一个8x8x25000阵列W和一个8 x 25000阵列r。我想通过r的每列(8x1)将每个8x8切片的W复用多个,并将结果保存在Wres中,这将最终成为8x25000矩阵。

我正在使用for循环完成此操作:

cmk -D <hostname>

但这很慢,我希望有更快的方法来实现这一目标。

有什么想法吗?

2 个答案:

答案 0 :(得分:3)

只要2个阵列共享相同的1轴长度,Matmul就可以传播。来自文档:

  

如果任一参数是N-D,则N> 2,它被视为驻留在最后两个索引中的一堆矩阵并相应地进行广播。

因此,您必须在matmul之前执行2次操作:

import numpy as np
a = np.random.rand(8,8,100)
b = np.random.rand(8, 100)
  1. 转置ab,以便第一个轴是100个切片
  2. b添加额外维度,以便b.shape = (100, 8, 1)
  3. 然后:

     at = a.transpose(2, 0, 1) # swap to shape 100, 8, 8
     bt = b.T[..., None] # swap to shape 100, 8, 1
     c = np.matmul(at, bt)
    

    c现在为100, 8, 1,重新塑造为8, 100

     c = np.squeeze(c).swapaxes(0, 1)
    

     c = np.squeeze(c).T
    

    最后,仅仅是为了方便的单行:

    c = np.squeeze(np.matmul(a.transpose(2, 0, 1), b.T[..., None])).T
    

答案 1 :(得分:2)

使用np.matmul的替代方法是np.einsum,可以在1个更短且可以说更可口的代码行中完成,而不需要方法链接。

示例数组:

np.random.seed(123)
w = np.random.rand(8,8,25000)
r = np.random.rand(8,25000)
wres = np.einsum('ijk,jk->ik',w,r)

# a quick check on result equivalency to your loop
print(np.allclose(np.matmul(w[:, :, 1], r[:, 1]), wres[:, 1]))
True

时间相当于@ Imanol的解决方案,所以请选择这两个。两者都比循环快30倍。在这里,由于数组的大小,einsum将具有竞争力。如果阵列大于这些,它可能会胜出,而对于较小的阵列则会输。有关详情,请参阅this讨论。

def solution1():
    return np.einsum('ijk,jk->ik',w,r)

def solution2():
    return np.squeeze(np.matmul(w.transpose(2, 0, 1), r.T[..., None])).T

def solution3():
    Wres = np.empty((8, 25000))
    for i in range(0,25000):
        Wres[:,i] = np.matmul(w[:,:,i],r[:,i])
    return Wres

%timeit solution1()
100 loops, best of 3: 2.51 ms per loop

%timeit solution2()
100 loops, best of 3: 2.52 ms per loop

%timeit solution3()
10 loops, best of 3: 64.2 ms per loop

Credit:@Divakar