PyTorch模型参数统计问题

PyTorch模型参数统计问题

使用torchsummary

使用torchsummary可以方便的统计模型的参数量

from torchsummary import summary
model = net().cuda()
input_shape = (1,256,256)
summary(model, input_shape)

即可打印出网络结构和参数统计
但是有一个问题：对于共享参数的模块(说白了就是重复调用的模块)会重复统计
这里有一篇知乎专栏文章很直观地分析了这个问题。

以后用torchsummary统计参数时如若遇到参数量爆炸的情况可以用下面的方法进行统计，不过两个还是结合着来用比较好，毕竟torchsummary还可以打印网络结构，比较直观。

def count_parameters(model):
	return sum(p.numel() for p in model.parameters() 
		if p.requires_grad)

看一下对比：