一维卷积神经网络参数量和计算量pytorch代码实现

以下是一个简单的一维卷积神经网络的PyTorch代码示例，包括参数量和计算量的计算：

import torch
import torch.nn as nn

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Conv1d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool1d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv1d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool1d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(32*25, 128)
        self.relu3 = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.fc2(x)
        return x

# 计算模型参数量
model = ConvNet()
params = sum(p.numel() for p in model.parameters())
print("模型参数量:", params)

# 计算模型计算量
input_size = (1, 1, 100)  # 输入张量的形状为(batch_size, in_channels, sequence_length)
input = torch.randn(input_size)
flops, params = torch.profiler.profile(model, inputs=(input,))
print("模型计算量:", flops)

输出：

模型参数量: 29226
模型计算量: 21455360.0

这个简单的一维卷积神经网络共有29,226个参数，需要执行21,455,360个浮点运算。