计算模型的计算复杂度GFLPS

计算模型的计算复杂度GFLPS

from thop import profile
def GFLOPS(model):
    inputdata = torch.randn(1, 3, 513, 513)  #这里可以修改
    if torch.cuda.is_available():
        model1 = model.cuda()
        inputdata = inputdata.cuda()
    flops, params = profile(model1, inputs=(inputdata, ))
    print("GFLOPs :{:.2f}, Params : {:.2f}".format(flops/1e9,params/1e6)) #flops单位G,para单位M
GFLOPS(model)

你可能感兴趣的:(笔记,pytorch,深度学习,神经网络)