PyTorch之计算模型推理时间

一、参考资料

如何测试模型的推理速度
Pytorch 测试模型的推理速度

二、计算PyTorch模型推理时间

1. 计算CPU推理时间

import torch
import torchvision
import time
import tqdm
from torchsummary import summarydef calcCPUTime():model = torchvision.models.resnet18()model.eval()# summary(model, input_size=(3, 224, 224), device="cpu")dummy_input = torch.randn(1, 3, 224, 224)num_iterations = 1000  # 迭代次数# 预热, GPU 平时可能为了节能而处于休眠状态, 因此需要预热print('warm up ...\n')with torch.no_grad():for _ in range(100):_ = model(dummy_input)print('testing ...\n')total_forward_time = 0.0  # 使用time来测试# 记录开始时间start_event = time.time()with torch.no_grad():for _ in tqdm.tqdm(range(num_iterations)):start_forward_time = time.time()_ = model(dummy_input)end_forward_time = time.time()forward_time = end_forward_time - start_forward_timetotal_forward_time += forward_time * 1000  # 转换为毫秒# 记录结束时间end_event = time.time()elapsed_time = (end_event - start_event)  # 转换为秒fps = num_iterations / elapsed_timeelapsed_time_ms = elapsed_time / (num_iterations * dummy_input.shape[0])avg_forward_time = total_forward_time / (num_iterations * dummy_input.shape[0])print(f"FPS: {fps}")print("elapsed_time_ms:", elapsed_time_ms * 1000)print(f"Avg Forward Time per Image: {avg_forward_time} ms")if __name__ == "__main__":calcCPUTime()

输出结果

warm up ...testing ...100%|██████████| 1000/1000 [00:09<00:00, 102.13it/s]
FPS: 102.11109490533485
elapsed_time_ms: 9.793255090713501
Avg Forward Time per Image: 9.777164697647095 ms

CPU资源占用情况

2. 计算GPU推理时间

方法一

import torch
import torchvision
import time
import tqdm
from torchsummary import summarydef calcGPUTime():model = torchvision.models.resnet18()model.cuda()model.eval()# summary(model, input_size=(3, 224, 224), device="cuda")dummy_input = torch.randn(1, 3, 224, 224).cuda()num_iterations = 1000  # 迭代次数# 预热, GPU 平时可能为了节能而处于休眠状态, 因此需要预热print('warm up ...\n')with torch.no_grad():for _ in range(100):_ = model(dummy_input)print('testing ...\n')total_forward_time = 0.0  # 使用time来测试# 记录开始时间start_event = time.time() * 1000with torch.no_grad():for _ in tqdm.tqdm(range(num_iterations)):start_forward_time = time.time()_ = model(dummy_input)end_forward_time = time.time()forward_time = end_forward_time - start_forward_timetotal_forward_time += forward_time * 1000  # 转换为毫秒# 记录结束时间end_event = time.time() * 1000elapsed_time = (end_event - start_event) / 1000.0  # 转换为秒fps = num_iterations / elapsed_timeelapsed_time_ms = elapsed_time / (num_iterations * dummy_input.shape[0])avg_forward_time = total_forward_time / (num_iterations * dummy_input.shape[0])print(f"FPS: {fps}")print("elapsed_time_ms:", elapsed_time_ms * 1000)print(f"Avg Forward Time per Image: {avg_forward_time} ms")if __name__ == "__main__":calcGPUTime()

输出结果

warm up ...testing ...100%|██████████| 1000/1000 [00:01<00:00, 727.79it/s]
FPS: 727.1527832145586
elapsed_time_ms: 1.375226806640625
Avg Forward Time per Image: 1.3709843158721924 ms

GPU资源占用情况

方法二

import torch
import torchvision
import numpy as np
import tqdm# TODO - 计算模型的推理时间
def calcGPUTime():device = 'cuda:0'model = torchvision.models.resnet18()model.to(device)model.eval()repetitions = 1000dummy_input = torch.rand(1, 3, 224, 224).to(device)# 预热, GPU 平时可能为了节能而处于休眠状态, 因此需要预热print('warm up ...\n')with torch.no_grad():for _ in range(100):_ = model(dummy_input)# synchronize 等待所有 GPU 任务处理完才返回 CPU 主线程torch.cuda.synchronize()# 设置用于测量时间的 cuda Event, 这是PyTorch 官方推荐的接口,理论上应该最靠谱starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)# 初始化一个时间容器timings = np.zeros((repetitions, 1))print('testing ...\n')with torch.no_grad():for rep in tqdm.tqdm(range(repetitions)):starter.record()_ = model(dummy_input)ender.record()torch.cuda.synchronize()  # 等待GPU任务完成curr_time = starter.elapsed_time(ender)  # 从 starter 到 ender 之间用时,单位为毫秒timings[rep] = curr_timeavg = timings.sum() / repetitionsprint('\navg={}\n'.format(avg))if __name__ == '__main__':calcGPUTime()

输出结果

warm up ...testing ...100%|██████████| 1000/1000 [00:01<00:00, 627.50it/s]avg=1.4300348817110062

GPU资源占用情况