深度学习基础之参数量(3)

一般的CNN网络的参数量估计代码

python 复制代码
class ResidualBlock(nn.Module):
    def __init__(self, in_planes, planes, norm_fn='group', stride=1):
        super(ResidualBlock, self).__init__()
        print(in_planes, planes, norm_fn, stride)

        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1)
        self.relu = nn.ReLU(inplace=True)

        num_groups = planes // 8

        if norm_fn == 'group':
            self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
            self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
            if not stride == 1:
                self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)

        elif norm_fn == 'batch':
            self.norm1 = nn.BatchNorm2d(planes)
            self.norm2 = nn.BatchNorm2d(planes)
            if not stride == 1:
                self.norm3 = nn.BatchNorm2d(planes)

        elif norm_fn == 'instance':
            self.norm1 = nn.InstanceNorm2d(planes)
            self.norm2 = nn.InstanceNorm2d(planes)
            if not stride == 1:
                self.norm3 = nn.InstanceNorm2d(planes)

        elif norm_fn == 'none':
            self.norm1 = nn.Sequential()
            self.norm2 = nn.Sequential()
            if not stride == 1:
                self.norm3 = nn.Sequential()

        if stride == 1:
            self.downsample = None

        else:
            self.downsample = nn.Sequential(
                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3)

    def forward(self, x):
        print(x.shape)
        #exit()
        y = x
        y = self.relu(self.norm1(self.conv1(y)))
        y = self.relu(self.norm2(self.conv2(y)))

        if self.downsample is not None:
            x = self.downsample(x)

        return self.relu(x + y)


R=ResidualBlock(384, 384, norm_fn='instance', stride=1)
summary(R.to("cuda" if torch.cuda.is_available() else "cpu"), (384, 32, 32))

transformer结构的参数量的估计结果

python 复制代码
import torch
import torch.nn as nn
from thop import profile
from torchsummary import summary

# 定义一个简单的Transformer模型
class Transformer(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_heads, num_layers):
        super(Transformer, self).__init__()
        self.embedding = nn.Embedding(input_dim, hidden_dim)
        self.transformer_layers = nn.Transformer(
            d_model=hidden_dim,
            nhead=num_heads,
            num_encoder_layers=num_layers,
            num_decoder_layers=num_layers
        )
        self.fc = nn.Linear(hidden_dim, input_dim)

    def forward(self, src, tgt):
        src = self.embedding(src)
        tgt = self.embedding(tgt)
        output = self.transformer_layers(src, tgt)
        output = self.fc(output)
        return output

# 创建Transformer模型实例
model2 = Transformer(input_dim=512, hidden_dim=512, num_heads=8, num_layers=6)

# 使用thop进行FLOPS估算
flops, params = profile(model2, inputs=(torch.randint(0, 512, (128,)), torch.randint(0, 512, (64,))))
print(f"FLOPS: {flops / 1e9} G FLOPS")  # 打印FLOPS,以十亿FLOPS(GFLOPS)为单位

# 计算参数量并打印
num_params = sum(p.numel() for p in model2.parameters() if p.requires_grad)
print(f"Total number of trainable parameters: {num_params}")
相关推荐
C7211BA2 小时前
基于网页的大语言模型聊天机器人
人工智能·语言模型·机器人
铭瑾熙2 小时前
深度学习之人脸检测
人工智能·深度学习
白光白光3 小时前
量子卷积神经网络
人工智能·神经网络·cnn
goomind4 小时前
深度学习实战人脸识别
深度学习·计算机视觉·人脸识别·pyqt5·facenet·mtcnn·人脸定位
陈苏同学5 小时前
机器翻译 & 数据集 (NLP基础 - 预处理 → tokenize → 词表 → 截断/填充 → 迭代器) + 代码实现 —— 笔记3.9《动手学深度学习》
人工智能·pytorch·笔记·python·深度学习·自然语言处理·机器翻译
狂放不羁霸5 小时前
组会 | 大语言模型 + LoRA
人工智能·语言模型·自然语言处理
sp_fyf_20245 小时前
【大语言模型】ACL2024论文-20 SCIMON:面向新颖性的科学启示机器优化
人工智能·深度学习·机器学习·语言模型·自然语言处理·数据挖掘
宋138102797205 小时前
SouVR Feedback force7 力反馈设备
人工智能·机器人·vr
我叫白小猿5 小时前
【大模型-智能体】AutoGen Studio测试和导出工作流程
人工智能·python·workflow·工作流·智能体·autogen
CopyLower5 小时前
AI赋能电商:智能购物推荐、会员分类与商品定价的创新探索
人工智能·分类·数据挖掘