模型运行速度FPS测试以及参数Parameter计算代码

模型FPS速度计算

import time
import torch
import numpy as np

net = build_your_net
net.eval()

# x是输入图片的大小
x = torch.zeros((1,3,H,W)).cuda()
t_all = []

for i in range(100):
    t1 = time.time()
    y = net(x)
    t2 = time.time()
    t_all.append(t2 - t1)

print('average time:', np.mean(t_all) / 1)
print('average fps:',1 / np.mean(t_all))

print('fastest time:', min(t_all) / 1)
print('fastest fps:',1 / min(t_all))

print('slowest time:', max(t_all) / 1)
print('slowest fps:',1 / max(t_all))

模型参数计算1

import numpy as np

def modelsize(model,type_size=4):
    para = sum([np.prod(list(p.size())) for p in model.parameters()])
    print('Model {} : params: {:4f}M'.format(model._get_name(), para * type_size / 1000 / 1000))

net = build_your_model
modelsize(net)

模型参数计算2

采用thop库，通过
pip install thop

from thop import profile
total_ops,total_params = profile(your_net, (data,), verbose = False)
print("| %.2f | %.2f " % (total_params/(1000*2),total_ops / (1000*3)))