Pytorch计算网络参数的两种方法

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(10, 20)  # 10个输入特征到20个输出特征的全连接层
        self.fc2 = nn.Linear(20, 30)  # 20个输入特征到30个输出特征的全连接层
        # 假设还有一个ReLU激活层，但它没有参数

    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)  # 激活层
        x = self.fc2(x)
        return x

# 实例化网络
net = SimpleNet()

# 计算总参数量
total_params = sum(p.numel() for p in net.parameters() if p.requires_grad)
print(f'Total number of parameters: {total_params}')

pip install torchsummary

import torch
from torchsummary import summary

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

    def forward(self, x):
        x = torch.nn.functional.relu(self.conv1(x))
        x = torch.nn.functional.relu(self.conv2(x))
        x = x.view(-1, 128 * 32 * 32)
        x = torch.nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = Net()

summary(model, (3, 32, 32))