前言
模型的推理速度测试有两种方式:一种是使用python的时间戳time函数来记录,另一种是使用Pytorch里的Event。同时,在进行GPU测试时,为减少冷启动的状态影响,可以先进行预热。
代码如下(示例):
if __name__ == "__main__":
model= BiSeNet(backbone='STDCNet813', n_classes=2, export=True)
model.cuda()
model.eval()
x = torch.randn(16, 3, 256, 256).cuda()
num_iterations = 500 # 迭代次数
# 预热阶段
for _ in range(30):
output = model(x)
# 正式测试阶段 (1)使用Event来测试
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
total_forward_time = 0.0 # 使用time来测试
start_event.record() # 记录开始时间
for _ in range(num_iterations):
start_forward_time = time.time()
output = model(x)
end_forward_time = time.time()
forward_time = end_forward_time - start_forward_time
total_forward_time += forward_time * 1000 # 转换为毫秒
end_event.record() # 记录结束时间
# 同步等待GPU操作完成
torch.cuda.synchronize()
elapsed_time = start_event.elapsed_time(end_event) / 1000.0 # 转换为秒
fps = num_iterations / elapsed_time
elapsed_time_ms = elapsed_time / (num_iterations * x.shape[0])
avg_forward_time = total_forward_time / (num_iterations * x.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")
注:
(1)在进行时间测量时,确保在记录结束事件(end_event.record())之后立即调用 torch.cuda.synchronize() 方法以同步等待 GPU 操作完成是非常重要的。这样可以确保所有前向传递操作都已经完成,避免了异步操作的影响.
(2)通常,Python 的时间戳time仅提供了秒级精度,而 torch.cuda.Event 可以提供更高的精度。