pythonpython-2.7scikit-learntimeit

How can I capture return value with Python timeit module?


Im running several machine learning algorithms with sklearn in a for loop and want to see how long each of them takes. The problem is I also need to return a value and DONT want to have to run it more than once because each algorithm takes so long. Is there a way to capture the return value 'clf' using python's timeit module or a similar one with a function like this...

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

when I call the function like this

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

P.S. I also dont want to set a global 'clf' because I might want to do multithreading or multiprocessing later.


Solution

  • The problem boils down to timeit._template_func not returning the function's return value:

    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
    

    We can bend timeit to our will with a bit of monkey-patching:

    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))
    

    returns

    (1.0010340213775635, 42)
    

    The first value is the timeit result (in seconds), the second value is the function's return value.

    Note that the monkey-patch above only affects the behavior of timeit when a callable is passed timeit.Timer. If you pass a string statement, then you'd have to (similarly) monkey-patch the timeit.template string.