python - 加速 Kronecker 产品 Numpy

Question

因此，我正在尝试计算两个任意维度矩阵的克罗内克积。 (为了示例，我使用了相同维度的方阵)

最初我尝试使用 kron :

a = np.random.random((60,60))
b = np.random.random((60,60))

start = time.time()
a = np.kron(a,b)
end = time.time()

Output: 0.160096406936645

为了尝试加快速度，我使用了 tensordot:

a = np.random.random((60,60))
b = np.random.random((60,60))

start = time.time()
a = np.tensordot(a,b,axes=0)
a = np.transpose(a,(0,2,1,3))
a = np.reshape(a,(3600,3600))
end = time.time()

Output: 0.11808371543884277

在网络上稍作搜索后，我发现(或者至少据我所知)numpy 在必须 reshape 已转置的张量时会生成一个额外的副本。

然后我尝试了以下(这段代码显然没有给出 a 和 b 的克罗内克积，但我只是作为测试这样做):

a = np.random.random((60,60))
b = np.random.random((60,60))

start = time.time()
a = np.tensordot(a,b,axes=0)
a = np.reshape(a,(3600,3600))
end = time.time()

Output: 0.052041053771972656

我的问题是:如何在不遇到与转置相关的问题的情况下计算克罗内克积？

我只是在寻找更快的速度，所以解决方案不必使用 tensordot。

编辑

我刚刚在这个堆栈帖子中找到:speeding up numpy kronecker products ，还有另一种方法:

a = np.random.random((60,60))
b = np.random.random((60,60))

c = a

start = time.time()
a = a[:,np.newaxis,:,np.newaxis]
a = a[:,np.newaxis,:,np.newaxis]*b[np.newaxis,:,np.newaxis,:]
a.shape = (3600,3600)
end = time.time()

test = np.kron(c,b)
print(np.array_equal(a,test))
print(end-start)


Output: True
0.05503702163696289

我仍然对您是否可以进一步加快此计算的问题感兴趣？

Answer 1

einsum 似乎有效:

>>> a = np.random.random((60,60))
>>> b = np.random.random((60,60))
>>> ab = np.kron(a,b)
>>> abe = np.einsum('ik,jl', a, b).reshape(3600,3600)
>>> (abe==ab).all()
True
>>> timeit(lambda: np.kron(a, b), number=10)
1.0697475590277463
>>> timeit(lambda: np.einsum('ik,jl', a, b).reshape(3600,3600), number=10)
0.42500176999601535

简单的广播甚至更快:

>>> abb = (a[:, None, :, None]*b[None, :, None, :]).reshape(3600,3600)
>>> (abb==ab).all()
True
>>> timeit(lambda:  (a[:, None, :, None]*b[None, :, None, :]).reshape(3600,3600), number=10)
0.28011218502069823

更新:使用 blas 和 cython 我们可以获得另一个适度 (30%) 的加速。自己决定是否值得麻烦。

[设置.py]

from distutils.core import setup
from Cython.Build import cythonize

setup(name='kronecker',
      ext_modules=cythonize("cythkrn.pyx"))

[cythkrn.pyx]

import cython
cimport scipy.linalg.cython_blas as blas
import numpy as np

@cython.boundscheck(False)
@cython.wraparound(False)
def kron(double[:, ::1] a, double[:, ::1] b):
    cdef int i = a.shape[0]
    cdef int j = a.shape[1]
    cdef int k = b.shape[0]
    cdef int l = b.shape[1]
    cdef int onei = 1
    cdef double oned = 1
    cdef int m, n
    result = np.zeros((i*k, j*l), float)
    cdef double[:, ::1] result_v = result
    for n in range(i):
        for m in range(k):
            blas.dger(&l, &j, &oned, &b[m, 0], &onei, &a[n, 0], &onei, &result_v[m+k*n, 0], &l)
    return result

构建运行 cython cythkrn.pyx 然后 python3 setup.py build。

>>> from timeit import timeit
>>> import cythkrn
>>> import numpy as np
>>> 
>>> a = np.random.random((60,60))
>>> b = np.random.random((60,60))
>>>
>>> np.all(cythkrn.kron(a, b)==np.kron(a, b))
True
>>> 
>>> timeit(lambda: cythkrn.kron(a, b), number=10)
0.18925874299020506