Pytorch查看模型结构

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

准备

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

pip install torchsummary

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

pip install torchinfo

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

使用方法

模型定义

假设我们有一个简单的AlexNet（GPT完成）

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)

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

model = AlexNet2D(num_class=10)

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

print(model)

效果：

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中传入实例化的模型和输入参数尺寸。

import torchsummary

....

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

Output:

----------------------------------------------------------------
        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：

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

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

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

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

torchinfo

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

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

效果

==========================================================================================      
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等其他递归模型
• 查看指定层数的尺寸