PyTorch 查看模型结构：输出张量维度、参数个数¶
使用pytorch-summary实现Keras中model.summary()的类似效果。
在自定义网络结构时，我们可以用print(model)来查看网络的基本信息，但只能看到有哪些层，每一层是什么（BatchNorm2d,、MaxPool2d,、AvgPool2d 等等），并不能看到每一层的输出张量的维数、每一层有多少个参数。
而理解每一层输出张量的维数对于理解网络结构是至关重要的。Keras中的model.summary()可以实现这一效果。但这在PyTorch中不能直接做到，可以用torchsummary包实现类似的效果。
定义模型¶
print(model)查看简单的模型结构¶
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)查看层类型、输出维度和参数个数¶
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
----------------------------------------------------------------
官方指南¶
安装¶
导入包¶
官方示例¶
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 。