使用torch.cuda.Event() 精确记录GPU上模型的推理时间

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)