为什么 numpy 矩阵乘法计算时间在 100x100 时增加一个数量级？

Question

当使用 numpy 计算A @ a其中A是随机 N × N 矩阵并且 a 是具有 N 个随机元素的向量时，计算时间在 N=100 时会跳跃一个数量级。这有什么特别的原因吗？作为比较，在CPU上使用torch进行相同的操作有更逐渐的增加

\n

用 python3.10 和 3.9 和 3.7 尝试过，具有相同的行为

\n

用于生成绘图的 numpy 部分的代码：

\n

import numpy as np\nfrom tqdm.notebook import tqdm\nimport pandas as pd\nimport time\nimport sys\n\ndef sym(A):\n    return .5 * (A + A.T)\n\nresults = []\nfor n in tqdm(range(2, 500)):\n    for trial_idx in range(10):\n        A = sym(np.random.randn(n, n))\n        a = np.random.randn(n)        \n        \n        t = time.time()\n        for i in range(1000):\n            A @ a\n        t = time.time() - t\n        results.append({\n            \'n\': n,\n            \'time\': t,\n            \'method\': \'numpy\',\n        })\nresults = pd.DataFrame(results)\n\nfrom matplotlib import pyplot as plt\nfig, ax = plt.subplots(1, 1)\nax.semilogy(results.n.unique(), results.groupby(\'n\').time.mean(), label="numpy")\nax.set_title(f\'A @ a timimgs (1000 times)\\nPython {sys.version.split(" ")[0]}\')\nax.legend()\nax.set_xlabel(\'n\')\nax.set_ylabel(\'avg. time\')\n

\n

更新

\n

添加

\n

import os\nos.environ["MKL_NUM_THREADS"] = "1" \nos.environ["NUMEXPR_NUM_THREADS"] = "1" \nos.environ["OMP_NUM_THREADS"] = "1" \n

\n

\xc3\xacmport numpy在给出更预期的输出之前，请参阅此答案以了解详细信息： /sf/answers/5226349481/

\n

Answer 1

numpy 在将大小为 100 或更大的矩阵相乘时尝试使用线程，并且线程乘法的默认 CBLAS 实现...次优，与 intel-MKL 或 ATLAS 等其他后端相反。

如果您使用本文中的答案强制它仅使用 1 个线程，您将获得一条连续的线来提高 numpy 性能。