import torch
time_pool = 0
for input_tensor,label in dataloader:
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record(stream=torch.cuda.current_stream())
output_tensor = model(input_tensor)
end.record(stream=torch.cuda.current_stream())
end.synchronize()
time_pool += start.elapsed_time(end)
使用torch.cuda.Event() 精确记录GPU上模型的推理时间
最新推荐文章于 2024-07-08 00:01:13 发布