计算 ~1m Hermitian 矩阵的光谱范数：`numpy.linalg.norm` 太慢了

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% 使用。考虑到这一点，我查看了以前的 Whosebug 问题：

why isn't numpy.mean multithreaded?
Why does multiprocessing use only a single core after I import numpy?
multithreaded blas in python/numpy（没有帮助）

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

我考虑过将我的数组拆分成更小的块，然后并行处理它们（可能在使用 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 函数中。

Answer 2

np.linalg.norm(A, ord=2) 通过使用 SVD 找到最大的奇异值来计算谱范数。但是，由于您的 8x8 子矩阵是 Hermitian，因此它们的最大奇异值将等于它们的绝对特征值的最大值 (see here)：

import numpy as np

def random_symmetric(N, k):
    A = np.random.randn(N, k, k)
    A += A.transpose(0, 2, 1)
    return A

N = 100000
k = 8
A = random_symmetric(N, k)

norm1 = np.abs(np.linalg.eigvalsh(A)).max(1)
norm2 = np.linalg.norm(A, ord=2, axis=(1, 2))

print(np.allclose(norm1, norm2))
# True

Hermitian 矩阵上的特征分解比 SVD 快很多：

In [1]: %%timeit A = random_symmetric(N, k)
np.linalg.norm(A, ord=2, axis=(1, 2))
   ....: 
1 loops, best of 3: 1.54 s per loop

In [2]: %%timeit A = random_symmetric(N, k)
np.abs(np.linalg.eigvalsh(A)).max(1)
   ....: 
1 loops, best of 3: 757 ms per loop

计算 ~1m Hermitian 矩阵的光谱范数：`numpy.linalg.norm` 太慢了

Computing the spectral norms of ~1m Hermitian matrices: `numpy.linalg.norm` is too slow

python

numpy

cython

linear-algebra

numba