加速 numpy.dot

Question

我有一个 numpy 脚本，它在以下代码中花费了大约 50% 的运行时间：

s = numpy.dot(v1, v1)

在哪里

v1 = v[1:]

而v 是存储在连续内存中的float64 的4000 元素一维ndarray（v.strides 是(8,)）。

有什么加快速度的建议吗？

edit 这是在 Intel 硬件上。这是我的numpy.show_config() 的输出：

atlas_threads_info:
    libraries = ['lapack', 'ptf77blas', 'ptcblas', 'atlas']
    library_dirs = ['/usr/local/atlas-3.9.16/lib']
    language = f77
    include_dirs = ['/usr/local/atlas-3.9.16/include']

blas_opt_info:
    libraries = ['ptf77blas', 'ptcblas', 'atlas']
    library_dirs = ['/usr/local/atlas-3.9.16/lib']
    define_macros = [('ATLAS_INFO', '"\\"3.9.16\\""')]
    language = c
    include_dirs = ['/usr/local/atlas-3.9.16/include']

atlas_blas_threads_info:
    libraries = ['ptf77blas', 'ptcblas', 'atlas']
    library_dirs = ['/usr/local/atlas-3.9.16/lib']
    language = c
    include_dirs = ['/usr/local/atlas-3.9.16/include']

lapack_opt_info:
    libraries = ['lapack', 'ptf77blas', 'ptcblas', 'atlas']
    library_dirs = ['/usr/local/atlas-3.9.16/lib']
    define_macros = [('ATLAS_INFO', '"\\"3.9.16\\""')]
    language = f77
    include_dirs = ['/usr/local/atlas-3.9.16/include']

lapack_mkl_info:
  NOT AVAILABLE

blas_mkl_info:
  NOT AVAILABLE

mkl_info:
  NOT AVAILABLE

Answer 1

也许罪魁祸首是复制传递给 dot 的数组。

正如 Sven 所说，dot 产品依赖于 BLAS 操作。这些操作需要以连续的 C 顺序存储的数组。如果传递给 dot 的两个数组都在 C_CONTIGUOUS 中，您应该会看到更好的性能。

当然，如果您传递给 dot 的两个数组确实是 1D (8,)，那么您应该看到 both C_CONTIGUOUS 和 F_CONTIGUOUS 标志设置为 True；但如果它们是 (1, 8)，那么您可以看到混合顺序。

>>> w = NP.random.randint(0, 10, 100).reshape(100, 1)
>>> w.flags
   C_CONTIGUOUS : True
   F_CONTIGUOUS : False
   OWNDATA : False
   WRITEABLE : True
   ALIGNED : True
   UPDATEIFCOPY : False

另一种选择：使用 BLAS 中的 _GEMM，它通过模块 scipy.linalg.fblas 公开。 （这两个数组 A 和 B 显然是 Fortran 顺序的，因为使用了 fblas。）

from scipy.linalg import fblas as FB
X = FB.dgemm(alpha=1., a=A, b=B, trans_b=True)

Answer 2

如果编译正确，numpy.dot 将使用多线程。确保它与顶部。我知道人们没有在 numpy w/atlas 中使用多线程的情况。此外，值得尝试使用针对 intel mkl 库编译的 numpy 版本。它们包括应该比英特尔硬件上的 atlas 更快的 blas 例程。你可以试试 enthought 的 python 发行版。包含所有这些，对拥有 edu 电子邮件帐户的人免费。

Answer 3

您的阵列不是很大，所以 ATLAS 可能没有做太多。以下 Fortran 程序的时间安排是什么？假设 ATLAS 没有做太多，这应该让您了解如果没有任何 python 开销，dot() 可能有多快。使用 gfortran -O3 我可以得到 5 +/- 0.5 us 的速度。

    program test

    real*8 :: x(4000), start, finish, s
    integer :: i, j
    integer,parameter :: jmax = 100000

    x(:) = 4.65
    s = 0.
    call cpu_time(start)
    do j=1,jmax
        s = s + dot_product(x, x)
    enddo
    call cpu_time(finish)
    print *, (finish-start)/jmax * 1.e6, s

    end program test

Answer 4

我唯一能想到的加快速度的方法是确保您的 NumPy 安装是针对优化的 BLAS 库（如 ATLAS）进行编译的。 numpy.dot() 是少数几个使用 BLAS 的 NumPy 函数之一。