Pytorch 测试模型的推理速度

# TODO - 计算模型的推理时间
def calcTime():import numpy as npfrom torchvision.models import resnet50import torchfrom torch.backends import cudnnimport tqdm'''  导入你的模型from module.amsnet import amsnet, anet, msnet, iresnet18, anet2, iresnet2, amsnet2from module.resnet import resnet18, resnet34from module.alexnet import AlexNetfrom module.vgg import vggfrom module.lenet import LeNetfrom module.googLenet import GoogLeNetfrom module.ivgg import iVGG'''cudnn.benchmark = Truedevice = 'cuda:0'model = anet().to(device)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))