为何Numpy中100×100矩阵与向量相乘耗时陡增一个数量级？

当用Numpy计算N×N随机对称矩阵A与N维随机向量a的乘积A @ a时，发现N=100时计算耗时陡增一个数量级，但在CPU上用PyTorch执行相同操作时，耗时增长却平缓得多。该现象在Python 3.7、3.9、3.10版本中均存在。

复现Numpy性能曲线的代码 #

import numpy as np
from tqdm.notebook import tqdm
import pandas as pd
import time
import sys

def sym(A):
    return .5 * (A + A.T)

results = []
for n in tqdm(range(2, 500)):
    for trial_idx in range(10):
        A = sym(np.random.randn(n, n))
        a = np.random.randn(n)
        
        t = time.time()
        for i in range(1000):
            A @ a
        t = time.time() - t
        results.append({
            'n': n,
            'time': t,
            'method': 'numpy',
        })
results = pd.DataFrame(results)

from matplotlib import pyplot as plt
fig, ax = plt.subplots(1, 1)
ax.semilogy(results.n.unique(), results.groupby('n').time.mean(), label="numpy")
ax.set_title(f'A @ a timimgs (1000 times)\nPython {sys.version.split(" ")[0]}')
ax.legend()
ax.set_xlabel('n')
ax.set_ylabel('avg. time')

问题解决方法 #

在导入numpy之前，添加以下代码限制线程数，即可得到符合预期的平滑性能表现：

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

内容的提问来源于stack exchange，提问作者Linus