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

Question

我在for循环中使用sklearn运行几个机器学习算法,并希望看到每个机器学习算法需要多长时间.问题是我还需要返回一个值,并且DONT想要多次运行它,因为每个算法都需要很长时间.有没有办法使用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)

PS我也不想设置全局'clf',因为我可能想稍后进行多线程或多处理.

Answer 1

问题归结为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结果(以秒为单位),第二个值是函数的返回值.

需要注意的是猴子补丁仅影响以上的行为timeit时,可调用的传递timeit.Timer.如果你传递一个字符串语句,那么你必须(类似地)猴子修补timeit.template字符串.

Answer 2

对于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 3

有趣的是,我也在做机器学习,并且有类似的要求;-)

我通过编写一个函数解决了以下问题:

• 运行你的功能
• 打印运行时间以及功能名称
• 返回结果

假设您想要时间:

clf = RandomForest(train_input, train_output)

然后做:

clf = time_fn( RandomForest, train_input, train_output )

Stdout将显示如下内容:

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