Pytorch查看模型结构

控制台输出深度学习模型各层结构与特征图尺寸。

准备

对于2D模型和部分3D模型可以使用torchsummary

1
pip install torchsummary

对于3D模型则推荐使用torchinfo

1
pip install torchinfo

torchinfo算是对torchsummary的完全重写版,更为稳定,实际使用中推荐优先选择前者。

使用方法

模型定义

假设我们有一个简单的AlexNet(GPT完成)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
import torch
import torch.nn as nn
import torch.nn.functional as F

class AlexNet2D(nn.Module):
def __init__(self, num_classes=1000):
super(AlexNet2D, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2), # (N, 3, 224, 224) → (N, 64, 55, 55)
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2), # (N, 64, 27, 27)

nn.Conv2d(64, 192, kernel_size=5, padding=2), # → (N, 192, 27, 27)
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2), # → (N, 192, 13, 13)

nn.Conv2d(192, 384, kernel_size=3, padding=1), # → (N, 384, 13, 13)
nn.ReLU(inplace=True),
nn.Conv2d(384, 256, kernel_size=3, padding=1), # → (N, 256, 13, 13)
nn.ReLU(inplace=True),
nn.Conv2d(256, 256, kernel_size=3, padding=1), # → (N, 256, 13, 13)
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2), # → (N, 256, 6, 6)
)
self.classifier = nn.Sequential(
nn.Dropout(),
nn.Linear(256 * 6 * 6, 4096),
nn.ReLU(inplace=True),
nn.Dropout(),
nn.Linear(4096, 4096),
nn.ReLU(inplace=True),
nn.Linear(4096, num_classes)
)

def forward(self, x):
x = self.features(x)
x = x.view(x.size(0), 256 * 6 * 6)
x = self.classifier(x)
return x

print(model)

这是最简单直接的查看模型的方式,在定义模型:

1
model = AlexNet2D(num_class=10)

之后,直接使用print方法将模型的结构打印出来,但是知会打印模型init的部分:

1
print(model)

效果:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
AlexNet2D(
(features): Sequential(
(0): Conv2d(3, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2))
(1): ReLU(inplace=True)
(2): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
(3): Conv2d(64, 192, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
(4): ReLU(inplace=True)
(5): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
(6): Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(7): ReLU(inplace=True)
(8): Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(9): ReLU(inplace=True)
(10): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(11): ReLU(inplace=True)
(12): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
)
(classifier): Sequential(
(0): Dropout(p=0.5, inplace=False)
(1): Linear(in_features=9216, out_features=4096, bias=True)
(2): ReLU(inplace=True)
(3): Dropout(p=0.5, inplace=False)
(4): Linear(in_features=4096, out_features=4096, bias=True)
(5): ReLU(inplace=True)
(6): Linear(in_features=4096, out_features=1000, bias=True)
)
)

torchsummary

torchsummary.summary中传入实例化的模型和输入参数尺寸。

1
2
3
4
5
import torchsummary

....

torchsummary.summary(model,(3,224,224))

Output:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 64, 55, 55] 23,296
ReLU-2 [-1, 64, 55, 55] 0
MaxPool2d-3 [-1, 64, 27, 27] 0
Conv2d-4 [-1, 192, 27, 27] 307,392
ReLU-5 [-1, 192, 27, 27] 0
MaxPool2d-6 [-1, 192, 13, 13] 0
Conv2d-7 [-1, 384, 13, 13] 663,936
ReLU-8 [-1, 384, 13, 13] 0
Conv2d-9 [-1, 256, 13, 13] 884,992
ReLU-10 [-1, 256, 13, 13] 0
Conv2d-11 [-1, 256, 13, 13] 590,080
ReLU-12 [-1, 256, 13, 13] 0
MaxPool2d-13 [-1, 256, 6, 6] 0
Dropout-14 [-1, 9216] 0
Linear-15 [-1, 4096] 37,752,832
ReLU-16 [-1, 4096] 0
Dropout-17 [-1, 4096] 0
Linear-18 [-1, 4096] 16,781,312
ReLU-19 [-1, 4096] 0
Linear-20 [-1, 1000] 4,097,000
================================================================
Total params: 61,100,840
Total params: 61,100,840
Trainable params: 61,100,840
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.57
Forward/backward pass size (MB): 8.31
Params size (MB): 233.08
Estimated Total Size (MB): 241.96
----------------------------------------------------------------

注意

如果将模型挂在到了GPU:

1
2
3
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = AlexNet2D(num_classes=1000).to(device)

summary的部分也需要将模型挂在到GPU上

1
torchsummary.summary(model, (3, 224, 224), device=str(device))

torchinfo

使用方法和torchsummary几乎一直,只是需要手动指定batch_size的尺寸。但是输出的结构更为清晰

1
summary(model, input_size=(1, 3, 224, 224), device="cuda")

效果

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
==========================================================================================      
Layer (type:depth-idx) Output Shape Param #
==========================================================================================
AlexNet2D [1, 1000] --
├─Sequential: 1-1 [1, 256, 6, 6] --
│ └─Conv2d: 2-1 [1, 64, 55, 55] 23,296
│ └─ReLU: 2-2 [1, 64, 55, 55] --
│ └─MaxPool2d: 2-3 [1, 64, 27, 27] --
│ └─Conv2d: 2-4 [1, 192, 27, 27] 307,392
│ └─ReLU: 2-5 [1, 192, 27, 27] --
│ └─MaxPool2d: 2-6 [1, 192, 13, 13] --
│ └─Conv2d: 2-7 [1, 384, 13, 13] 663,936
│ └─ReLU: 2-8 [1, 384, 13, 13] --
│ └─Conv2d: 2-9 [1, 256, 13, 13] 884,992
│ └─ReLU: 2-10 [1, 256, 13, 13] --
│ └─Conv2d: 2-11 [1, 256, 13, 13] 590,080
│ └─ReLU: 2-12 [1, 256, 13, 13] --
│ └─MaxPool2d: 2-13 [1, 256, 6, 6] --
├─Sequential: 1-2 [1, 1000] --
│ └─Dropout: 2-14 [1, 9216] --
│ └─Linear: 2-15 [1, 4096] 37,752,832
│ └─ReLU: 2-16 [1, 4096] --
│ └─Dropout: 2-17 [1, 4096] --
│ └─Linear: 2-18 [1, 4096] 16,781,312
│ └─ReLU: 2-19 [1, 4096] --
│ └─Linear: 2-20 [1, 1000] 4,097,000
==========================================================================================
Total params: 61,100,840
Trainable params: 61,100,840
Non-trainable params: 0
Total mult-adds (M): 714.68
==========================================================================================
Input size (MB): 0.60
Forward/backward pass size (MB): 3.95
Params size (MB): 244.40
Estimated Total Size (MB): 248.96
==========================================================================================

3D的情况下

使用方法一致,只是需要注意输入数据的形状。

补充torchinfo

  • 支持RNNs、LSTM等其他递归模型
  • 查看指定层数的尺寸
作者

Zhou

发布于

2025-07-03

更新于

2025-07-03

许可协议

评论

+ + +