统计深度学习模型的参数和显存占用（pytorch）

最新推荐文章于 2024-04-11 17:48:42 发布

cactus233

最新推荐文章于 2024-04-11 17:48:42 发布

阅读量5.3k

点赞数 8

分类专栏：深度学习

本文链接：https://blog.csdn.net/zhaoxr233/article/details/103667640

版权

深度学习专栏收录该内容

2 篇文章 0 订阅

订阅专栏

这里介绍两种方法

１、使用相关的库torchsummary

from torchsummary import summary

net=net.to(torch.device("cpu"))#or cuda
summary(net,(4,228,912),device="cpu") #or cuda

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

2、使用库thop

from thop import profile
net=net.cuda()
input= torch.ones([1,4,128,128]).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))

这个比较简单，最后简单的输出参数量和其中的浮点计算次数