显示进度

Show timeit progress

我有多个函数,我反复想测量使用内置 timeit 库的执行时间。 (假设 fun1、fun2 和 fun3 都依赖于几个子例程,其中一些我正在尝试优化。每次迭代后,我想知道我的 3 个顶级函数的执行速度)

问题是,我事先不确定这些功能会持续多长时间 运行,我只是粗略估计。使用 timeit.repeat(...) 和足够数量的 repetitions/number 执行给了我一个很好的估计,但有时它需要很长时间,因为我不小心减慢了其中一个子例程。为计时例程设置一个类似 tqdm 的进度条会非常方便,这样我就可以提前估计在计时完成之前我必须等待多长时间。我在 timeit 库中没有找到任何这样的功能,所以这里是问题:

使用 timeit.repeat 或 timeit.timeit 计时功能时是否可以显示(类似 tqdm 的)进度条?

查看 source code of timeit,有一个模板会在任何计时结束时执行。可以简单地更改该模板以包含进度指示器:

import timeit

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

# some timeit test:
timeit.timeit(lambda: "-".join(map(str, range(100))), number=1000000)

当然,这会影响结果,因为 tqdm-调用在 _t0_t1 测量中。不过,tqdm 的文档声称每次迭代的开销仅为 60ns。

您可以创建 timeit.Timer 的子类,它使用 tqdm 来跟踪执行的总迭代次数。

from timeit import Timer, default_number
from tqdm import tqdm
import itertools
import gc

class ProgressTimer(Timer):
    def timeit(self, number=default_number):
        """Time 'number' executions of the main statement.
        To be precise, this executes the setup statement once, and
        then returns the time it takes to execute the main statement
        a number of times, as a float measured in seconds.  The
        argument is the number of times through the loop, defaulting
        to one million.  The main statement, the setup statement and
        the timer function to be used are passed to the constructor.
        """
        # wrap the iterator in tqdm
        it = tqdm(itertools.repeat(None, number), total=number)
        gcold = gc.isenabled()
        gc.disable()
        try:
            timing = self.inner(it, self.timer)
        finally:
            if gcold:
                gc.enable()
        # the tqdm bar sometimes doesn't flush on short timers, so print an empty line
        print()
        return timing

要使用这个对象,我们只需要传入我们想要的脚本运行。您可以将其定义为字符串(如下所示),也可以简单地打开文件进行读取并读取到变量。

py_setup = 'import numpy as np'

py_script = """
x = np.random.rand(1000)
x.sum()
"""

pt = ProgressTimer(py_script, setup=py_setup)
pt.timeit()

# prints / returns:
100%|███████████████████████████████████████████████| 1000000/1000000 [00:13<00:00, 76749.68it/s]
13.02982600001269