python - 显示时间进度

Question

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

问题是，我事先不确定这些函数要运行多长时间，我只是有一个粗略的估计。使用 timeit.repeat(...) 进行足够的重复/执行次数可以给我一个很好的估计，但有时需要很长时间，因为我不小心减慢了其中一个子例程的速度。如果计时例程有一个类似 tqdm 的进度条会非常方便，这样我就可以提前估计需要等待计时完成的时间。我在 timeit 库中没有找到任何此类功能，所以问题是:

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

Answer 1

您可以创建 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