python的时间测试小工具：timeit

    timeit是Python标准库内置的小工具，可以快速测试小段代码的性能。本文主要记录它的参数接口和使用示例

一、函数接口说明

        如下是timeit有两个函数接口：timeit.timeit, timeit.repeat。区别就是time.repeat可以定义重复的次数，返回的是重复每个number次的时间数组；timeit.timeit只返回代码段执行1次number次的时间。

default_number = 1000000

default_repeat = 5

def timeit(stmt="pass", setup="pass", timer=default_timer,
           number=default_number, globals=None):
    """Convenience function to create Timer object and call timeit method."""
    return Timer(stmt, setup, timer, globals).timeit(number)

def repeat(stmt="pass", setup="pass", timer=default_timer,
           repeat=default_repeat, number=default_number, globals=None):
    """Convenience function to create Timer object and call repeat method."""
    return Timer(stmt, setup, timer, globals).repeat(repeat, number)

以timeit 函数为例进行接口参数含义说明：

timeit.timeit(stmt, setup,timer, number)

参数说明：

• stmt: statement的缩写，你要测试的代码或者语句，纯文本，默认值是 "pass"
• setup: 在运行stmt前的配置语句，纯文本，默认值也是 "pass"
• timer: 计时器，一般忽略这个参数
• number: stmt执行的次数，默认是1000000，一百万

repeat 函数多一个repeat参数，repeat默认5次

timeit.repeat(stmt, setup, timer, repeat, number)

如下是同一个代码段，执行两个函数的返回区别

import timeit
print(timeit.timeit('output = 10*5')) # 0.06458526299866207
print(timeit.repeat('output = 10*5')) # [0.04510639099862601, 0.016582202000790858, 0.016219472001466784, 0.016255639999144478, 0.016456094001114252]

二、参数的填写拓展

  测试多行代码可以用分号来连接语句。

print(timeit.timeit(stmt='a=10;b=10;sum=a+b'))

也可以用三引号来写stmt。

import timeit
import_module = "import random"
testcode = ''' 
def test(): 
    return random.randint(10, 100)
'''
print(timeit.repeat(stmt=testcode, setup=import_module))

 

三、基于模块测试的示例

可以用来对比测试某个封装好的函数的性能

如下，写了几种matmul功能的函数接口，用timeit.timeit做性能测试

import numpy as np
import timeit

def native_matmul(arrA, arrB, arrC, M, N, K):
    for m in range(M):
        for n in range(N):
            tmp = 0
            for k in range(K):
                tmp += arrA[m,k] * arrB[k, n]
            arrC[m,n]=tmp

def loop_matmul(arrA, arrB, arrC, M, N, K):
    for m in range(M):
        tmp = np.zeros([N])
        for k in range(K):
            for n in range(N):
                tmp[n] += arrA[m,k] * arrB[k, n]
        arrC[m,:]=tmp

def loop1_matmul(arrA, arrB, arrC, M, N, K):
    for m in range(M):
        for k in range(K):
            for n in range(N):
                arrC[m,n] += arrA[m,k] * arrB[k, n]


def np_matmul(arrA, arrB, arrC):
    np.dot(arrA, arrB, arrC)



if __name__ =="__main__":

    test_context ='''
import numpy as np
from __main__ import native_matmul,loop_matmul,loop1_matmul,np_matmul
M = 128
N = 128
K = 128
arrA = np.random.random([M, K])
arrB = np.random.random([K, N])
arrC = np.zeros([M, N])
'''
    number = 10
    native_matmul_running_time = timeit.timeit(setup=test_context,
                                        stmt='native_matmul(arrA, arrB, arrC, M, N, K)',
                                        number=number, )
    loop_matmul_running_time = timeit.timeit(setup=test_context,
                                        stmt='loop_matmul(arrA, arrB, arrC, M, N, K)',
                                        number=number, )
    loop1_matmul_running_time = timeit.timeit(setup=test_context,
                                         stmt='loop1_matmul(arrA, arrB, arrC, M, N, K)',
                                         number=number, )
    np_matmul_running_time = timeit.timeit(setup=test_context,
                                         stmt='np_matmul(arrA, arrB, arrC)',
                                         number=number, )


    print("native_matmul_running_time", native_matmul_running_time/number)
    print("loop_matmul_running_time", loop_matmul_running_time/number)
    print("loop1_matmul_running_time", loop1_matmul_running_time/number)
    print("np_matmul_running_time", np_matmul_running_time/number)

函数输出：

native_matmul_running_time 0.882451231900086
loop_matmul_running_time 1.2186987614999452
loop1_matmul_running_time 1.3516903318999538
np_matmul_running_time 0.0001239164999788045

参考：使用timeit测试Python函数的性能 - 知乎