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

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

深度学习

PyTorch

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

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

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

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

定义模型¶

Pythontest2 = Inception2(combination, combination_rev, index_list_2)

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

Pythonprint(test2)

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

Text OnlyInception2(

(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)查看层类型、输出维度和参数个数¶

Pythonfrom 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

----------------------------------------------------------------

官方指南¶

安装¶

Bashpip install torchsummary

导入包¶

Pythonfrom torchsummary import summary

官方示例¶

Pythonimport 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 。

2024 年 3 月 23 日

评论