python-3.x - 如何使用 timeit 对函数进行计时并保存结果

Question

我正在尝试使用 timeit 比较我的函数的速度，但保存这些结果却无法实现。让我向您展示一个玩具示例，尝试比较不同的矩阵乘法函数:

from timeit import timeit
import numpy as np

def prod_tensorial (A,B)  : return np.tensordot(A,B,(-1,0))
def prod_punto (A,B)      : return np.dot(A,B)
def prod_sumeinsten (A,B) : return np.einsum('ij,jk->ik',A,B)

A = np.random.rand(1000,2000)
B = np.random.rand(2000,3000)

fs = ( prod_tensorial, prod_punto, prod_sumeinsten )
ts = [ ]
ps = [ ]

for f in fs:
    ts += [ timeit('ps += [f(A,B)]',number=1) ]     # error: ps & f didn't recognized
    #ts += [ timeit(lambda: f(A,B),number=1) ]      # it works but lossing results
    print( ts[-1] )

print(np.allclose(*ps))

如你所见，我无法保存结果。你知道我该怎么做吗？

Answer 1

函数timeit无法直接访问全局变量。您必须通过 globals 关键字提供全局变量的字典。由于变量 ps 在定时语句中首先读取，因此默认情况下将其视为本地变量。将其标记为全局:

# This part is only for testing
ps = []
def f(x,y): return x+y
A, B = 10, 11

timeit('global ps; ps += [f(A,B)]', number=1, globals=globals())