这篇文章主要介绍了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}
本文以 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
关于“Pytorch怎么统计参数网络参数数量”这篇文章的内容就介绍到这里,感谢各位的阅读!相信大家对“Pytorch怎么统计参数网络参数数量”知识都有一定的了解,大家如果还想学习更多知识,欢迎关注亿速云行业资讯频道。
免责声明:本站发布的内容(图片、视频和文字)以原创、转载和分享为主,文章观点不代表本网站立场,如果涉及侵权请联系站长邮箱:is@yisu.com进行举报,并提供相关证据,一经查实,将立刻删除涉嫌侵权内容。