跳转至

PyTorch 查看模型结构:输出张量维度、参数个数

使用pytorch-summary实现Kerasmodel.summary()的类似效果。

在自定义网络结构时,我们可以用print(model)来查看网络的基本信息,但只能看到有哪些层,每一层是什么(BatchNorm2d,、MaxPool2d,、AvgPool2d 等等),并不能看到每一层的输出张量的维数、每一层有多少个参数。

而理解每一层输出张量的维数对于理解网络结构是至关重要的。Keras中的model.summary()可以实现这一效果。但这在PyTorch中不能直接做到,可以用torchsummary包实现类似的效果。

定义模型

Python
test2 = Inception2(combination, combination_rev, index_list_2)

print(model)查看简单的模型结构

Python
print(test2)

print(model)只能查看简单的模型结构,不能看出每一层的输出张量的维数、每一层有多少个参数。

Text Only
Inception2(
  (bc1): BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (bc2): BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (bc3): BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (bc4): BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (bc5): BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (bc6): BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (bc7): BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (max_pool): MaxPool2d(kernel_size=(1, 10), stride=(1, 10), padding=0, dilation=1, ceil_mode=False)
  (avg_pool): AvgPool2d(kernel_size=(1, 10), stride=(1, 10), padding=0)
  (min_pool): MaxPool2d(kernel_size=(1, 10), stride=(1, 10), padding=0, dilation=1, ceil_mode=False)
  (bc_pool1): BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (bc_pool2): BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (bc_pool3): BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)

torchsummary.summary(model, input_size)查看层类型、输出维度和参数个数

Python
from torchsummary import summary

summary(test2, input_size=(1, 9, 30))
Text Only
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
       BatchNorm2d-1            [-1, 1, 36, 10]               2
       BatchNorm2d-2            [-1, 1, 36, 10]               2
       BatchNorm2d-3             [-1, 1, 9, 10]               2
       BatchNorm2d-4             [-1, 1, 9, 10]               2
       BatchNorm2d-5             [-1, 1, 9, 10]               2
       BatchNorm2d-6             [-1, 1, 9, 10]               2
       BatchNorm2d-7             [-1, 1, 9, 10]               2
         MaxPool2d-8            [-1, 1, 117, 1]               0
       BatchNorm2d-9            [-1, 1, 117, 1]               2
        AvgPool2d-10            [-1, 1, 117, 1]               0
      BatchNorm2d-11            [-1, 1, 117, 1]               2
        MaxPool2d-12            [-1, 1, 117, 1]               0
      BatchNorm2d-13            [-1, 1, 117, 1]               2
================================================================
Total params: 20
Trainable params: 20
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.01
Params size (MB): 0.00
Estimated Total Size (MB): 0.02
----------------------------------------------------------------

官方指南

安装

Bash
pip install torchsummary

导入包

Python
from torchsummary import summary

官方示例

Python
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") # PyTorch v0.4.0
model = Net().to(device)

summary(model, (1, 28, 28))
Text Only
----------------------------------------------------------------
        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
----------------------------------------------------------------

更多指南可参考https://github.com/sksq96/pytorch-summary

评论