如何使用 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())

如何使用 timeit 对函数计时并保存结果

How to time functions with timeit and save results

timeit

python-3.x