非Jupyter/iPython环境下%timeit的普通Python替代方案

你猜测的方向完全正确，Python标准库自带的timeit就是IPython/Jupyter中%timeit魔法命令的底层依赖，完全可以实现一致的计时效果，在Spyder、Thonny这类常规IDE的普通Python运行环境里都能直接用。

%timeit的核心逻辑是自动调整代码运行次数、多轮重复测试取最优平均耗时，尽可能降低系统调度、后台进程带来的计时误差，以下是几种常用的等效写法：

1. 基础调用：快速对比代码性能 #

直接调用timeit.timeit()即可完成基础计时，适合对比小段代码的运行速度，比如对比NumPy和JAX矩阵运算的场景：

import timeit
import numpy as np
import jax.numpy as jnp

# 测试原生NumPy矩阵乘法耗时
numpy_cost = timeit.timeit(
    stmt="np.dot(arr1, arr2)",
    setup="import numpy as np; arr1 = np.random.rand(1000,1000); arr2 = np.random.rand(1000,1000)",
    number=100  # 指定重复运行100次，返回总耗时
)
print(f"NumPy单次平均耗时：{numpy_cost/100:.6f}s")

# 测试JAX矩阵乘法耗时，注意JAX有JIT编译、异步调度特性，必须做预热+同步等待
jax_setup_code = """
import jax.numpy as jnp
arr1 = jnp.random.rand(1000,1000)
arr2 = jnp.random.rand(1000,1000)
# 预热：触发一次JIT编译，避免把编译时间算入运算耗时
res = jnp.dot(arr1, arr2)
res.block_until_ready()
"""
jax_cost = timeit.timeit(
    stmt="res = jnp.dot(arr1, arr2); res.block_until_ready()",
    setup=jax_setup_code,
    number=100
)
print(f"JAX单次平均耗时：{jax_cost/100:.6f}s")

注意：测试JAX、PyTorch这类带异步执行、JIT编译机制的框架时，必须调用同步接口（JAX对应block_until_ready()，PyTorch对应torch.cuda.synchronize()）等待计算实际完成，同时提前做一次预热运行，否则计时结果会完全失真，这一点和Jupyter环境下使用%timeit的注意事项完全一致。

2. 完全复刻%timeit的使用体验 #

如果想要和Jupyter里%timeit完全一致的效果——自动根据代码运行速度选择合适的循环次数（快代码多跑降误差，慢代码少跑省时间）、多轮测试取最优值、自动匹配时间单位输出，可以封装一个简单的工具函数：

import timeit

def magic_timeit(code: str, setup: str = "", repeat_round: int = 7, globals_env=None):
    """
    等效IPython %timeit的计时函数
    :param code: 要计时的代码字符串
    :param setup: 初始化/预热代码，正式计时前只运行一次
    :param repeat_round: 重复测试轮数，默认和%timeit一致为7轮
    :param globals_env: 传入全局变量环境，方便直接使用当前上下文里的变量
    """
    if globals_env is None:
        globals_env = globals()
    timer = timeit.Timer(stmt=code, setup=setup, globals=globals_env)
    # 自动选择最优的单轮循环次数
    loop_num, _ = timer.autorange()
    # 跑多轮取结果
    round_costs = timer.repeat(repeat=repeat_round, number=loop_num)
    # 和%timeit逻辑一致，取所有轮次里最快的单次平均耗时
    best_cost = min(round_costs) / loop_num

    # 自动适配时间单位，和%timeit输出格式对齐
    if best_cost < 1e-6:
        print_val = best_cost * 1e9
        unit = "ns"
    elif best_cost < 1e-3:
        print_val = best_cost * 1e6
        unit = "µs"
    elif best_cost < 1:
        print_val = best_cost * 1e3
        unit = "ms"
    else:
        print_val = best_cost
        unit = "s"
    print(f"{loop_num} loops, best of {repeat_round}: {print_val:.2f} {unit} per loop")
    return best_cost

使用时直接传入要测试的代码即可，输出效果和%timeit完全一致：

magic_timeit(
    code="res = jnp.dot(arr1, arr2); res.block_until_ready()",
    setup=jax_setup_code
)

3. 长代码块临时计时方案 #

如果要测试的代码逻辑比较长，不方便拆成字符串传入timeit，可以直接用time.perf_counter()做手动计时，适合调试阶段快速看耗时：

import time

# 先做预热
res = jnp.dot(arr1, arr2)
res.block_until_ready()

# 正式计时
run_times = 100
start_ts = time.perf_counter()
for _ in range(run_times):
    # 这里放你要计时的任意代码，不用转成字符串
    res = jnp.dot(arr1, arr2)
    res.block_until_ready()
end_ts = time.perf_counter()
print(f"单次平均耗时：{(end_ts - start_ts)/run_times:.6f}s")

注意：不要用time.time()做性能计时，它是系统墙钟时间，会受系统时间调整影响，精度远低于time.perf_counter()，timeit模块底层默认使用的也是perf_counter作为计时时钟。

内容的提问来源于stack exchange，提问作者Bastiaan Quast