Pytorch怎么统计参数网络参数数量

这篇文章主要介绍了Pytorch怎么统计参数网络参数数量的相关知识，内容详细易懂，操作简单快捷，具有一定借鉴价值，相信大家阅读完这篇Pytorch怎么统计参数网络参数数量文章都会有所收获，下面我们一起来看看吧。

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 torchimport torch.nn as nnclass 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 outclass 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 output3def 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 Networkin_channel = 128net1 = 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

关于“Pytorch怎么统计参数网络参数数量”这篇文章的内容就介绍到这里，感谢各位的阅读！相信大家对“Pytorch怎么统计参数网络参数数量”知识都有一定的了解，大家如果还想学习更多知识，欢迎关注编程网行业资讯频道。