计算深度学习模型的推理时间、fps

测量PyTorch模型推理速度的方法

最新推荐文章于 2024-09-15 17:06:55 发布

原创

最新推荐文章于 2024-09-15 17:06:55 发布 · 7.7k 阅读

43 ·

CC 4.0 BY-SA版权

文章标签：

#python #深度学习

这篇博客介绍了如何正确地衡量PyTorch模型的推理时间。关键在于使用torch.cuda.synchronize()确保GPU操作完成后再进行时间计算。提供的代码示例展示了如何通过measure_inference_speed函数来基准测试模型的性能，平均计算每秒处理的图像数量和每张图像所需的时间。

背景

代码

参考：

背景

评估模型的推理时间时有需要注意的地方。如torch.cuda.synchronize()，因为pytorch代码执行时异步的，使用该代码会等待gpu上所有操作结束后再接着运行代码、计算时间等【1】。

代码

函数【2】：


import time
def measure_inference_speed(model, data, max_iter=200, log_interval=50):
    model.eval()

    # the first several iterations may be very slow so skip them
    num_warmup = 5
    pure_inf_time = 0
    fps = 0

    # benchmark with 2000 image and take the average
    for i in range(max_iter):

        torch.cuda.synchronize()
        start_time = time.perf_counter()

        with torch.no_grad():
            model(*data)

        torch.cuda.synchronize()
        elapsed = time.perf_counter() - start_time

        if i >= num_warmup:
            pure_inf_time += elapsed
            if (i + 1) % log_interval == 0:
                fps = (i + 1 - num_warmup) / pure_inf_time
                print(
                    f'Done image [{i