MATLAB 矩阵乘法性能比 NumPy 快 5 倍

Question

我在 MATLAB 和 Python 中设置了两个关于广播矩阵乘法的相同测试。对于 Python，我使用了 NumPy，对于 MATLAB，我使用了使用 BLAS 的 mtimesx 库。

MATLAB

close all; clear;

N = 1000 + 100; % a few initial runs to be trimmed off at the end

a = 100;
b = 30;
c = 40;
d = 50;
A = rand(b, c, a);
B = rand(c, d, a);
C = zeros(b, d, a);

times = zeros(1, N);
for ii = 1:N
    tic
    C = mtimesx(A,B);
    times(ii) = toc;
end

times = times(101:end) * 1e3;

plot(times);
grid on;
title(median(times));

Python

import timeit
import numpy as np
import matplotlib.pyplot as plt


N = 1000 + 100  # a few initial runs to be trimmed off at the end

a = 100
b = 30
c = 40
d = 50
A = np.arange(a * b * c).reshape([a, b, c])
B = np.arange(a * c * d).reshape([a, c, d])
C = np.empty(a * b * d).reshape([a, b, d])

times = np.empty(N)

for i in range(N):
    start = timeit.default_timer()
    C = A @ B
    times[i] = timeit.default_timer() - start

times = times[101:] * 1e3

plt.plot(times, linewidth=0.5)
plt.grid()
plt.title(np.median(times))
plt.show()

• 我的 Python 是从 pip 安装的默认 Python，它使用 OpenBLAS。
• 我在英特尔 NUC i3 上运行。

MATLAB 代码在 1ms 内运行，而 Python 在 5.8ms 内运行，我不知道为什么，因为它们似乎都在使用 BLAS。

---

编辑

来自 Anaconda：

In [7]: np.__config__.show()
mkl_info:
    libraries = ['mkl_rt']
    library_dirs = [...]
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = [...]
blas_mkl_info:
    libraries = ['mkl_rt']
    library_dirs = [...]
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = [...]
blas_opt_info:
    libraries = ['mkl_rt']
    library_dirs = [...]
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = [...]
lapack_mkl_info:
    libraries = ['mkl_rt']
    library_dirs = [...]
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = [...]
lapack_opt_info:
    libraries = ['mkl_rt']
    library_dirs = [...]
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = [...]

从 numpy 使用 pip

In [2]: np.__config__.show()
blas_mkl_info:
NOT AVAILABLE
blis_info:
NOT AVAILABLE
openblas_info:
    library_dirs = [...]
    libraries = ['openblas']
    language = f77
    define_macros = [('HAVE_CBLAS', None)]
blas_opt_info:
    library_dirs = [...]
    libraries = ['openblas']
    language = f77
    define_macros = [('HAVE_CBLAS', None)]
lapack_mkl_info:
NOT AVAILABLE
openblas_lapack_info:
    library_dirs = [...]
    libraries = ['openblas']
    language = f77
    define_macros = [('HAVE_CBLAS', None)]
lapack_opt_info:
    library_dirs = [...]
    libraries = ['openblas']
    language = f77
    define_macros = [('HAVE_CBLAS', None)]

---

编辑 2 我尝试用 np.matmul(A, B, out=C) 替换 C = A @ B 并得到 2x worse 时间，例如大约 11 毫秒。这真的很奇怪。

Answer 1

您的 MATLAB 代码使用浮点数组，但 NumPy 代码使用整数数组。这在时间上有很大的不同。对于 MATLAB 和 NumPy 之间的“苹果对苹果”比较，Python/NumPy 代码还必须使用浮点数组。

然而，这并不是唯一重要的问题。 NumPy github 站点的issue 7569（以及issue 8957）中讨论的 NumPy 存在缺陷。 “堆叠”数组的矩阵乘法不使用快速 BLAS 例程来执行乘法。这意味着多维数组的乘法可能比预期的要慢得多。

二维数组的乘法确实使用快速例程，因此您可以通过在循环中将各个二维数组相乘来解决此问题。令人惊讶的是，尽管 Python 循环的开销很大，但它在许多情况下比应用于完整堆叠数组的 @、matmul 或 einsum 更快。

这是 NumPy 问题中显示的函数的变体，它在 Python 循环中执行矩阵乘法：

def xmul(A, B):
    """
    Multiply stacked matrices A (with shape (s, m, n)) by stacked
    matrices B (with shape (s, n, p)) to produce an array with
    shape (s, m, p).

    Mathematically equivalent to A @ B, but faster in many cases.

    The arguments are not validated.  The code assumes that A and B
    are numpy arrays with the same data type and with shapes described
    above.
    """
    out = np.empty((a.shape[0], a.shape[1], b.shape[2]), dtype=a.dtype)
    for j in range(a.shape[0]):
        np.matmul(a[j], b[j], out=out[j])
    return out

我的 NumPy 安装也使用 MKL（它是 Anaconda 发行版的一部分）。这是A @ B 和xmul(A, B) 的时序比较，使用浮点值数组：

In [204]: A = np.random.rand(100, 30, 40)

In [205]: B = np.random.rand(100, 40, 50)

In [206]: %timeit A @ B
4.76 ms ± 6.37 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [207]: %timeit xmul(A, B)
582 µs ± 35.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

尽管xmul 使用了 Python 循环，但它所花费的时间大约是 A @ B 的 1/8。

Answer 2

我认为这是内存排序的问题。 Matlab的zeros(a, b, c)和numpy的zeros((a, b, c), order='F')一样，不是默认的。

当然，正如您已经确定的那样，@ 在与mtimesx 不同的轴上运行。为了使比较公平，您应该确保您的数组在 matlab 顺序中，然后转置以处理语义差异

# note: `order` in reshape actually changes the resulting array data,
# not just its memory layout
A = np.arange(a * b * c).reshape([b, c, a], order='F').transpose((2, 0, 1))
B = np.arange(a * c * d).reshape([c, d, a], order='F').transpose((2, 0, 1))

Answer 3

您能用最近发布的 NumPy 1.16 再试一次吗？我们对 matmul 进行了重构，将 BLAS 用于内部二维，这应该会加快代码速度。