Pytorch统计参数网络参数数量
def get_parameter_number(net):
total_num = sum(p.numel() for p in net.parameters())
trainable_num = sum(p.numel() for p in net.parameters() if p.requires_grad)
return {'Total': total_num, 'Trainable': trainable_num}
Pytorch如何计算网络的参数量
本文以 Dense Block 为例，Pytorch 为 DL 框架，最终计算模块参数量方法如下：
import torch
import torch.nn as nn
class Norm_Conv(nn.Module):
def __init__(self,in_channel):
super(Norm_Conv,self).__init__()
self.layers = nn.Sequential(
nn.Conv2d(in_channel,in_channel,3,1,1),
nn.ReLU(True),
nn.BatchNorm2d(in_channel),
nn.Conv2d(in_channel,in_channel,3,1,1),
nn.ReLU(True),
nn.BatchNorm2d(in_channel),
nn.Conv2d(in_channel,in_channel,3,1,1),
nn.ReLU(True),
nn.BatchNorm2d(in_channel))
def forward(self,input):
out = self.layers(input)
return out
class DenseBlock_Norm(nn.Module):
def __init__(self,in_channel):
super(DenseBlock_Norm,self).__init__()
self.first_layer = nn.Sequential(nn.Conv2d(in_channel,in_channel,3,1,1),
nn.ReLU(True),
nn.BatchNorm2d(in_channel))
self.second_layer = nn.Sequential(nn.Conv2d(in_channel*2,in_channel,3,1,1),
nn.ReLU(True),
nn.BatchNorm2d(in_channel))
self.third_layer = nn.Sequential(
nn.Conv2d(in_channel*3,in_channel,3,1,1),
nn.ReLU(True),
nn.BatchNorm2d(in_channel))
def forward(self,input):
output1 = self.first_layer(input)
output2 = self.second_layer(torch.cat((output1,input),dim=1))
output3 = self.third_layer(torch.cat((input,output1,output2),dim=1))
return output3
def count_param(model):
param_count = 0
for param in model.parameters():
param_count += param.view(-1).size()[0]
return param_count
# Get Parameter number of Network
in_channel = 128
net1 = Norm_Conv(in_channel)
print('Norm Conv parameter count is {}'.format(count_param(net1)))
net2 = DenseBlock_Norm(in_channel)
print('DenseBlock Norm parameter count is {}'.format(count_param(net2)))最终结果如下
Norm Conv parameter count is 443520
DenseBlock Norm parameter count is 885888
总结
以上为个人经验，希望能给大家一个参考，也希望大家多多支持编程网。
