pytorch计算模型参数量

PyTorch是一个流行的深度学习框架，可以用于构建和训练神经网络模型。计算模型参数量是评估模型复杂度和性能的重要指标之一。在PyTorch中，可以使用`torchsummary`库来计算模型参数量。

首先，确保你已经安装了`torchsummary`库。可以使用以下命令进行安装：

pip install torchsummary

接下来，我们可以通过以下步骤来计算模型参数量：

1. 导入所需的库：

import torch
from torchsummary import summary

2. 定义模型：

import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(128 * 7 * 7, 256)
        self.fc2 = nn.Linear(256, 10)

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

model = MyModel()

这是一个简单的卷积神经网络模型。

3. 使用`summary`函数计算模型参数量：

summary(model, (3, 224, 224))

这里的`(3, 224, 224)`是输入张量的形状。

运行以上代码后，你将会得到一个模型参数量的总结，包括每一层的输出形状和参数量。