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}

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)))