深度学习模型的参数和显存占用计算

１、使用相关的库torchsummary

参数量、浮点数计算量、中间变量、ｔｒａｉｎ的变量数、保持不变的变量数，每一层的中间变量和类型都会详细列出

from torchsummary import summary
net=net.to(torch.device("cpu"))
summary(net,(3,224,224),device="cpu") 

2、使用库thop

简单的输出参数量和其中的浮点计算次数

from thop import profile
net=net.cuda()
input= torch.ones([1,3,224,224]).cuda()
inputs=[]
inputs.append(input)
flops, params=profile(net,inputs)#,custom_ops={model.ResNet,countModel})
print("flops:{0:,} ".format(flops))
print("parms:{0:,}".format(params))