PyTorch中如何查看神经网络模型的参数（两种高效的方法，简单上手）

1 用for循环打印parameters

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=3)
        self.fc1 = nn.Linear(1350, 10)

    def forward(self, x):
        print(x.size())  # 结果：[1, 1, 32, 32]
        # 卷积 -> 激活 -> 池化
        x = self.conv1(x)
        x = F.relu(x)
        print(x.size())  # 结果：[1, 6, 30, 30]
        x = F.max_pool2d(x, (2, 2))
        x = F.relu(x)
        print(x.size())  # 结果：[1, 6, 15, 15]
        # reshape，‘-1’表示自适应
        # 这里做的就是压扁的操作 就是把后面的[1, 6, 15, 15]压扁，变为 [1, 1350]
        x = x.view(x.size()[0], -1)
        print(x.size())  # 这里就是fc1层的的输入1350
        x = self.fc1(x)
        return x

• 如下，先实例化Net，然后打印输出

net = Net()

for parameters in net.parameters():
    print(parameters)

print("################################################")
for name, parameters in net.named_parameters():
    print(name, ':', parameters.size())

• 输出结果：

Parameter containing:
tensor([[[[ 0.0102,  0.0416,  0.0778],
          [-0.0645, -0.0658,  0.0764],
          [-0.1447, -0.0363,  0.0263]]],


        [[[-0.1047,  0.2682,  0.1654],
          [-0.2595, -0.0150, -0.0464],
          [-0.3228, -0.3133, -0.2081]]],


        [[[-0.2690,  0.1592, -0.1979],
          [ 0.3166,  0.2864, -0.1338],
          [ 0.3009, -0.1523,  0.1127]]],


        [[[ 0.0622,  0.0731, -0.0764],
          [-0.1877, -0.1849,  0.2464],
          [ 0.3300,  0.2713,  0.3299]]],


        [[[ 0.0554,  0.2507, -0.3236],
          [ 0.1371, -0.0487,  0.2763],
          [-0.1178,  0.1598, -0.2895]]],


        [[[ 0.2391,  0.3132, -0.3278],
          [-0.2279, -0.2666,  0.1694],
          [-0.0407, -0.3296,  0.2220]]]], requires_grad=True)
Parameter containing:
tensor([ 0.3329,  0.1968, -0.0663,  0.2801, -0.1373, -0.0967],
       requires_grad=True)
Parameter containing:
tensor([[ 0.0048,  0.0091,  0.0096,  ..., -0.0250, -0.0161,  0.0105],
        [ 0.0236, -0.0030,  0.0263,  ..., -0.0109,  0.0207, -0.0005],
        [-0.0230,  0.0205,  0.0266,  ...,  0.0227,  0.0004, -0.0221],
        ...,
        [ 0.0058,  0.0036, -0.0029,  ...,  0.0249, -0.0157, -0.0078],
        [-0.0221, -0.0050, -0.0080,  ..., -0.0153,  0.0051,  0.0119],
        [-0.0011,  0.0062,  0.0176,  ..., -0.0143, -0.0010, -0.0078]],
       requires_grad=True)
Parameter containing:
tensor([ 0.0091, -0.0270,  0.0096, -0.0145,  0.0180,  0.0224, -0.0099, -0.0205,
         0.0018, -0.0199], requires_grad=True)
################################################
conv1.weight : torch.Size([6, 1, 3, 3])
conv1.bias : torch.Size([6])
fc1.weight : torch.Size([10, 1350])
fc1.bias : torch.Size([10])

2 安装依赖：torchsummary

1. 首先安装

pip torchsummary

2. 导包+调用（注意传参的结构）

from torchsummary import summary
summary(model, input_size=(channels, H, W))  # 这里是函数原型

2.1 如果是单输入，比如CNN 模型

2.1.1 代码1

import torch

import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torchsummary import summary

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=3)
        self.fc1 = nn.Linear(1350, 10)

    def forward(self, x):
        print(x.size())  # 结果：[1, 1, 32, 32]
        # 卷积 -> 激活 -> 池化
        x = self.conv1(x)
        x = F.relu(x)
        print(x.size())  # 结果：[1, 6, 30, 30]
        x = F.max_pool2d(x, (2, 2))
        x = F.relu(x)
        print(x.size())  # 结果：[1, 6, 15, 15]
        # reshape，‘-1’表示自适应
        # 这里做的就是压扁的操作 就是把后面的[1, 6, 15, 15]压扁，变为 [1, 1350]
        x = x.view(x.size()[0], -1)
        print(x.size())  # 这里就是fc1层的的输入1350
        x = self.fc1(x)
        return x


net = Net()

for parameters in net.parameters():
    print(parameters)

print("################################################")
for name, parameters in net.named_parameters():
    print(name, ':', parameters.size())

print("################################################")
summary(net, input_size=(1, 32, 32))

• 输出：【重点看最后一个summary的输出，之前的输出同上】

torch.Size([2, 1, 32, 32])
torch.Size([2, 6, 30, 30])
torch.Size([2, 6, 15, 15])
torch.Size([2, 1350])
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1            [-1, 6, 30, 30]              60
            Linear-2                   [-1, 10]          13,510
================================================================
Total params: 13,570
Trainable params: 13,570
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.04
Params size (MB): 0.05
Estimated Total Size (MB): 0.10
----------------------------------------------------------------

Process finished with exit code 0

2.1.2 代码2

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Net().to(device)

summary(model, (1, 28, 28))

• 输出情况：

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 10, 24, 24]             260
            Conv2d-2             [-1, 20, 8, 8]           5,020
         Dropout2d-3             [-1, 20, 8, 8]               0
            Linear-4                   [-1, 50]          16,050
            Linear-5                   [-1, 10]             510
================================================================
Total params: 21,840
Trainable params: 21,840
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.06
Params size (MB): 0.08
Estimated Total Size (MB): 0.15
----------------------------------------------------------------

Process finished with exit code 0

2.2 如果是多输入的情况，比如说RNN需要输入x和h_0

• 只要传入的input_size改为一个安装输入所需size组成的列表就行

import torch
import torch.nn as nn
from torchsummary import summary


class SimpleConv(nn.Module):
    def __init__(self):
        super(SimpleConv, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
        )

    def forward(self, x, y):
        x1 = self.features(x)
        x2 = self.features(y)
        return x1, x2


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleConv().to(device)

summary(model, input_size=[(1, 16, 16), (1, 28, 28)])

• 输出：

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1            [-1, 1, 16, 16]              10
              ReLU-2            [-1, 1, 16, 16]               0
            Conv2d-3            [-1, 1, 28, 28]              10
              ReLU-4            [-1, 1, 28, 28]               0
================================================================
Total params: 20
Trainable params: 20
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.77
Forward/backward pass size (MB): 0.02
Params size (MB): 0.00
Estimated Total Size (MB): 0.78
----------------------------------------------------------------

Process finished with exit code 0