Pytorch 中怎么计算网络的参数量

最新推荐文章于 2024-03-29 17:55:31 发布

小张Python1

最新推荐文章于 2024-03-29 17:55:31 发布

阅读量3.4k

点赞数 3

分类专栏： Pytorch 文章标签： Pytorch DenseBlock Norm_Conv 参数量计算

本文链接：https://blog.csdn.net/weixin_42512684/article/details/115674336

版权

本博文以 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),