使numpy.linalg.svd和numpy矩阵乘法使用多线程

Question

我有一个脚本，它使用了大量的numpy和numpy.linalg函数，经过一些研究之后，我发现它们应该会自动使用多线程。当然，我的htop显示总是只显示一个线程来运行我的脚本。

我对多线程很陌生，现在还不太清楚如何正确设置它。

我主要是利用numpy.linalg.svd

这是numpy.show_config()的输出

openblas64__info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None)]
    runtime_library_dirs = ['/usr/local/lib']
blas_ilp64_opt_info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None)]
    runtime_library_dirs = ['/usr/local/lib']
openblas64__lapack_info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None), ('HAVE_LAPACKE', None)]
    runtime_library_dirs = ['/usr/local/lib']
lapack_ilp64_opt_info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None), ('HAVE_LAPACKE', None)]
    runtime_library_dirs = ['/usr/local/lib']
Supported SIMD extensions in this NumPy install:
    baseline = SSE,SSE2,SSE3
    found = SSSE3,SSE41,POPCNT,SSE42,AVX,F16C,FMA3,AVX2
    not found = AVX512F,AVX512CD,AVX512_KNL,AVX512_KNM,AVX512_SKX,AVX512_CLX,AVX512_CNL,AVX512_ICL

MRE

import numpy as np 
import tensorly as ty 


tensor = np.random.rand(32,32,32,32)

unfolding = ty.unfold(tensor,0)
unfolding = unfolding @ unfolding.transpose()
U,S,_ = np.linalg.svd(unfolding)

更新

正如公认的答案所建议的那样，用MKL重建numpy解决了这个问题。

Answer 1

主要问题是矩阵的大小太小，对于线程来说在所有平台上都是不值得的。实际上，OpenBLAS使用OpenMP创建有关矩阵大小的线程。线程通常只创建一次，但对于机器(通常是普通PC上的数百微秒)，创建时间从几十微秒到几十毫秒不等。机器上的核心数量越多，要创建的线程数量就越多，因此开销也就越大。当OpenMP线程池被重用时，仍然需要支付开销，主要是由于工作的分布和线程之间的同步，尽管开销通常要小得多(通常是一个数量级的小)。

尽管如此，当输出矩阵比输入矩阵小时，OpenBLAS做出了明显的次优选择--(这就是你的情况)。实际上，OpenBLAS在运行目标内核之前很难知道并行开销，所以它必须做出选择:通常根据输入矩阵的大小设置一个阈值，以便定义内核将在何时顺序执行或使用多个线程。这对于非常小的内核仍然是快速的，对于保持与其他BLAS实现竞争的巨大内核来说，这是至关重要的。问题是，这个门槛并没有被完美的选择。它看起来像OpenBLAS只看输出矩阵的大小，这显然是次优的“薄”矩阵，如在您的代码(例如。50x1000000 @ 1000000x50)经验分析表明，在您的情况下，阈值是任意设置为100x100的:超过这个阈值，OpenBLAS使用多个线程，而不是其他线程。问题是，在大多数平台上，线程对于小得多的矩阵已经很有用了。对于64x64x64x64张量)。

这个阈值是由编译时定义()调优的，如GEMM_MULTITHREAD_THRESHOLD (用于gemm.c (或gemv.c ))。注意，在代码中，k维很重要，但这并不是基准测试在我的机器上所显示的(可能是由于使用了较旧版本的OpenBLAS )。您可以使用更小的阈值(比如1而不是4)重新构建OpenBLAS。

另一种解决方案是使用另一个BLAS实现，如BLIS或Intel MKL，它们应该使用不同的阈值(可能更好)。最后一种解决方案是实现特定于的实现，以高效地计算代码的矩阵(可能使用Numba或Cython)，但是BLAS实现经过了很大的优化，因此通常很难实际编写更快的代码(除非您非常熟悉低级别的优化、编译器和现代处理器体系结构)。