python - 计算 ~1m Hermitian 矩阵的光谱范数 : `numpy.linalg.norm` is too slow

Question

我想计算 N 8x8 Hermitian 矩阵的谱范数，N 接近 1E6。例如，以这 100 万个随机复数 8x8 矩阵为例:

import numpy as np

array = np.random.rand(8,8,1e6)  + 1j*np.random.rand(8,8,1e6)

目前使用 numpy.linalg.norm 需要我将近 10 秒:

np.linalg.norm(array, ord=2, axis=(0,1))

我尝试使用下面的 Cython 代码，但这只给我带来了微不足道的性能改进:

import numpy as np
cimport numpy as np
cimport cython

np.import_array()

DTYPE = np.complex64

@cython.boundscheck(False)
@cython.wraparound(False)
def function(np.ndarray[np.complex64_t, ndim=3] Array):
    assert Array.dtype == DTYPE
    cdef int shape0 = Array.shape[2]
    cdef np.ndarray[np.float32_t, ndim=1] normarray = np.zeros(shape0, dtype=np.float32)
    normarray = np.linalg.norm(Array, ord=2, axis=(0, 1))
    return normarray

我还尝试了 numba 和其他一些 scipy 函数(例如 scipy.linalg.svdvals)来计算这些矩阵的奇异值。一切还是太慢了。

有没有可能让它变得更快？ numpy 是否已经优化到使用 Cython 或 numba 无法提高速度的程度？还是我的代码效率很低，我做的事情从根本上是错误的？

我注意到在进行计算时，我的 CPU 内核中只有两个被 100% 使用。考虑到这一点，我查看了以前的 StackOverflow 问题:

• why isn't numpy.mean multithreaded?

• Why does multiprocessing use only a single core after I import numpy?

• multithreaded blas in python/numpy (没有帮助)

和其他几个人，但不幸的是我仍然没有解决方案。

我考虑过将我的数组拆分成更小的 block ，然后并行处理它们(可能在使用 CUDA 的 GPU 上)。在 numpy/Python 中有没有办法做到这一点？我还不知道我的代码中的瓶颈在哪里，即它是 CPU 还是内存限制，或者可能是其他原因。

Answer 1

深入研究 np.linalg.norm 的代码，我推断，对于这些参数，它正在寻找 N 维上矩阵奇异值的最大值

首先生成一个小样本数组。使 N 成为消除 rollaxis 操作的第一个维度:

In [268]: N=10; A1 = np.random.rand(N,8,8)+1j*np.random.rand(N,8,8)

In [269]: np.linalg.norm(A1,ord=2,axis=(1,2))
Out[269]: 
array([ 5.87718306,  5.54662999,  6.15018125,  5.869058  ,  5.80882818,
        5.86060462,  6.04997992,  5.85681085,  5.71243196,  5.58533323])

等价操作:

In [270]: np.amax(np.linalg.svd(A1,compute_uv=0),axis=-1)
Out[270]: 
array([ 5.87718306,  5.54662999,  6.15018125,  5.869058  ,  5.80882818,
        5.86060462,  6.04997992,  5.85681085,  5.71243196,  5.58533323])

相同的值，相同的时间:

In [271]: timeit np.linalg.norm(A1,ord=2,axis=(1,2))
1000 loops, best of 3: 398 µs per loop
In [272]: timeit np.amax(np.linalg.svd(A1,compute_uv=0),axis=-1)
1000 loops, best of 3: 389 µs per loop

而大部分时间花在svd上，产生一个(N,8)数组:

In [273]: timeit np.linalg.svd(A1,compute_uv=0)
1000 loops, best of 3: 366 µs per loop

因此，如果您想加速 norm，您需要进一步研究加速此 svd。 svd 使用 np.linalg._umath_linalg 函数 - 这是一个 .so 文件 - 已编译。

c代码在https://github.com/numpy/numpy/blob/97c35365beda55c6dead8c50df785eb857f843f0/numpy/linalg/umath_linalg.c.src中

看起来这确实是您将获得的最快速度。没有 Python 级循环。任何循环都在 c 代码或它调用的 lapack 函数中。