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 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怎么統計參數網絡參數數量”知識都有一定的了解，大家如果還想學習更多知識，歡迎關注億速云行業資訊頻道。