如何用Python timeit模块捕获返回值？

Question

我正在用sklearn在一个循环里运行几个机器学习算法，想要看看每个算法花了多长时间。问题是我还需要返回一个值，而且不想多跑一次，因为每个算法都很耗时。有没有办法用Python的timeit模块或者类似的工具来捕捉返回值'clf'，像这样调用函数...

def RandomForest(train_input, train_output):
    clf = ensemble.RandomForestClassifier(n_estimators=10)
    clf.fit(train_input, train_output)
    return clf

当我这样调用这个函数时

t = Timer(lambda : RandomForest(trainX,trainy))
print t.timeit(number=1)

另外，我也不想设置一个全局的'clf'，因为我可能以后想要做多线程或者多进程处理。

Answer 1

截至2020年，在ipython或jupyter notebook中，它是这样的：

t = %timeit -n1 -r1 -o RandomForest(trainX, trainy)
t.best

Answer 2

如果我理解得没错，从 Python 3.5 开始，你可以在每个 Timer 实例中定义全局变量，而不需要在你的代码块里先定义它们。我不太确定这样做在并行处理时是否会遇到同样的问题。

我的做法大概是这样的：

clf = ensemble.RandomForestClassifier(n_estimators=10)
myGlobals = globals()
myGlobals.update({'clf'=clf})
t = Timer(stmt='clf.fit(trainX,trainy)', globals=myGlobals)
print(t.timeit(number=1))
print(clf)

Answer 3

有趣的是，我也在做机器学习，遇到了类似的需求 ;-)

我通过写一个函数来解决这个问题，这个函数：

• 运行你的函数
• 打印出运行时间，以及你函数的名字
• 返回结果

假设你想要计时：

clf = RandomForest(train_input, train_output)

那么你可以这样做：

clf = time_fn( RandomForest, train_input, train_output )

输出会显示类似这样的内容：

mymodule.RandomForest: 0.421609s

下面是 time_fn 的代码：

import time

def time_fn( fn, *args, **kwargs ):
    start = time.clock()
    results = fn( *args, **kwargs )
    end = time.clock()
    fn_name = fn.__module__ + "." + fn.__name__
    print fn_name + ": " + str(end-start) + "s"
    return results

Answer 4

对于Python 3.5，你可以覆盖 timeit.template 的值。

timeit.template = """
def inner(_it, _timer{init}):
    {setup}
    _t0 = _timer()
    for _i in _it:
        retval = {stmt}
    _t1 = _timer()
    return _t1 - _t0, retval
"""

unutbu的回答 在Python 3.4中有效，但在3.5中无效，因为_template_func这个函数似乎在3.5中被移除了。

Answer 5

这个问题的关键在于 timeit._template_func 没有返回函数的返回值：

def _template_func(setup, func):
    """Create a timer function. Used if the "statement" is a callable."""
    def inner(_it, _timer, _func=func):
        setup()
        _t0 = _timer()
        for _i in _it:
            _func()
        _t1 = _timer()
        return _t1 - _t0
    return inner

我们可以通过一些小技巧来调整 timeit 的行为：

import timeit
import time

def _template_func(setup, func):
    """Create a timer function. Used if the "statement" is a callable."""
    def inner(_it, _timer, _func=func):
        setup()
        _t0 = _timer()
        for _i in _it:
            retval = _func()
        _t1 = _timer()
        return _t1 - _t0, retval
    return inner

timeit._template_func = _template_func

def foo():
    time.sleep(1)
    return 42

t = timeit.Timer(foo)
print(t.timeit(number=1))

返回值

(1.0010340213775635, 42)

第一个值是 timeit 的结果（以秒为单位），第二个值是函数的返回值。

需要注意的是，上面的调整只会影响当你把一个可调用对象（callable）传给 timeit.Timer 时的行为。如果你传的是一个字符串语句，那么你也需要对 timeit.template 的字符串进行类似的调整。