使用 PyTorch 深度学习搭建模型后，如果想查看模型结构，可以直接使用 print(model) 函数打印。但该输出结果不是特别直观，查阅发现有个能输出类似 keras 风格 model.summary() 的模型可视化工具。这里记录一下方便以后查阅。
PyTorch 打印模型结构、输出维度和参数信息（torchsummary）
安装 torchsummary
pip install torchsummary
输出网络信息
summary函数介绍
model：网络模型
input_size：网络输入图片的shape，这里不用加batch_size进去
batch_size：batch_size参数，默认是-1
device：在GPU还是CPU上运行，默认是cuda在GPU上运行，如果想在CPU上执行将参数改为CPU即可
import torch
import torch.nn as nn
from torchsummary import summary
class Shallow_ConvNet(nn.Module):
def __init__(self, in_channel, conv_channel_temp, kernel_size_temp, conv_channel_spat, kernel_size_spat,
pooling_size, pool_stride_size, dropoutRate, n_classes, class_kernel_size) :
super(Shallow_ConvNet, self).__init__()
self.temp_conv = nn.Conv2d(in_channels=in_channel,
out_channels=conv_channel_temp,
kernel_size=(1, kernel_size_temp),
stride=1,
bias=False)
self.spat_conv = nn.Conv2d(in_channels=conv_channel_temp,
out_channels=conv_channel_spat,
kernel_size=(kernel_size_spat, 1),
stride=1,
bias=False)
self.bn = nn.BatchNorm2d(num_features=conv_channel_spat)
# slef.act_conv = x*x
self.pooling = nn.AvgPool2d(kernel_size=(1, pooling_size),
stride=(1, pool_stride_size))
# slef.act_pool = log(max(x, eps))
self.dropout = nn.Dropout(p=dropoutRate)
self.class_conv = nn.Conv2d(in_channels=conv_channel_spat,
out_channels=n_classes,
kernel_size=(1, class_kernel_size),
bias=False)
self.softmax = nn.Softmax(dim=1)
def safe_log(self, x):
""" Prevents :math:`log(0)` by using :math:`log(max(x, eps))`."""
return torch.log(torch.clamp(x, min=1e-6))
def forward(self, x):
# input shape (batch_size, C, T)
if len(x.shape) is not 4:
x = torch.unsqueeze(x, 1)
# input shape (batch_size, 1, C, T)
x = self.temp_conv(x)
x = self.spat_conv(x)
x = self.bn(x)
x = x*x # conv_activate
x = self.pooling(x)
x = self.safe_log(x) # pool_activate
x = self.dropout(x)
x = self.class_conv(x)
x= self.softmax(x)
out = torch.squeeze(x)
return out
###============================ Initialization parameters ============================###
channels = 44
samples = 534
in_channel = 1
conv_channel_temp = 40
kernel_size_temp = 25
conv_channel_spat = 40
kernel_size_spat = channels
pooling_size = 75
pool_stride_size = 15
dropoutRate = 0.3
n_classes = 4
class_kernel_size = 30
def main():
input = torch.randn(32, 1, channels, samples)
model = Shallow_ConvNet(in_channel, conv_channel_temp, kernel_size_temp, conv_channel_spat, kernel_size_spat,
pooling_size, pool_stride_size, dropoutRate, n_classes, class_kernel_size)
out = model(input)
print('===============================================================')
print('out', out.shape)
print('model', model)
summary(model=model, input_size=(1,channels,samples), batch_size=32, device="cpu")
if __name__ == "__main__":
main()
输出：
out torch.Size([32, 4])
model Shallow_ConvNet(
(temp_conv): Conv2d(1, 40, kernel_size=(1, 25), stride=(1, 1), bias=False)
(spat_conv): Conv2d(40, 40, kernel_size=(44, 1), stride=(1, 1), bias=False)
(bn): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(pooling): AvgPool2d(kernel_size=(1, 75), stride=(1, 15), padding=0)
(dropout): Dropout(p=0.3, inplace=False)
(class_conv): Conv2d(40, 4, kernel_size=(1, 30), stride=(1, 1), bias=False)
(softmax): Softmax(dim=1)
)
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [32, 40, 44, 510] 1,000
Conv2d-2 [32, 40, 1, 510] 70,400
BatchNorm2d-3 [32, 40, 1, 510] 80
AvgPool2d-4 [32, 40, 1, 30] 0
Dropout-5 [32, 40, 1, 30] 0
Conv2d-6 [32, 4, 1, 1] 4,800
Softmax-7 [32, 4, 1, 1] 0
================================================================
Total params: 76,280
Trainable params: 76,280
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 2.87
Forward/backward pass size (MB): 229.69
Params size (MB): 0.29
Estimated Total Size (MB): 232.85
----------------------------------------------------------------
AttributeError: ‘tuple’ object has no attribute ‘size’
旧的summary加入LSTM之类的会报错，需要用新的summarry
pip install torchinfo
from torchinfo import summary
def main():
input = torch.randn(32, window_size, channels, samples)
model = Cascade_Conv_LSTM(in_channel, out_channel_conv1, out_channel_conv2, out_channel_conv3, kernel_conv123, stride_conv123, padding_conv123,
fc1_in, fc1_out, dropoutRate1, lstm1_in, lstm1_hidden, lstm1_layer, lstm2_in, lstm2_hidden, lstm2_layer, fc2_in, fc2_out, dropoutRate2,
fc3_in, n_classes)
# model = model.to('cuda:1')
# input = torch.from_numpy(input).to('cuda:1').to(torch.float32).requires_grad_()
out = model(input)
print('===============================================================')
print('out', out.shape)
print('model', model)
summary(model=model, input_size=(32,10,channels,samples), device="cpu")
if __name__ == "__main__":
main()
==========================================================================================
Layer (type:depth-idx) Output Shape Param #
==========================================================================================
Cascade_Conv_LSTM [32, 4] --
├─Sequential: 1-1 [320, 32, 10, 11] --
│ └─Conv2d: 2-1 [320, 32, 10, 11] 288
│ └─ELU: 2-2 [320, 32, 10, 11] --
├─Sequential: 1-2 [320, 64, 10, 11] --
│ └─Conv2d: 2-3 [320, 64, 10, 11] 18,432
│ └─ELU: 2-4 [320, 64, 10, 11] --
├─Sequential: 1-3 [320, 128, 10, 11] --
│ └─Conv2d: 2-5 [320, 128, 10, 11] 73,728
│ └─ELU: 2-6 [320, 128, 10, 11] --
├─Sequential: 1-4 [320, 1024] --
│ └─Linear: 2-7 [320, 1024] 14,418,944
│ └─ELU: 2-8 [320, 1024] --
├─Dropout: 1-5 [320, 1024] --
├─LSTM: 1-6 [32, 10, 1024] 8,396,800
├─LSTM: 1-7 [32, 10, 1024] 8,396,800
├─Sequential: 1-8 [32, 1024] --
│ └─Linear: 2-9 [32, 1024] 1,049,600
│ └─ELU: 2-10 [32, 1024] --
├─Dropout: 1-9 [32, 1024] --
├─Linear: 1-10 [32, 4] 4,100
├─Softmax: 1-11 [32, 4] --
==========================================================================================
Total params: 32,358,692
Trainable params: 32,358,692
Non-trainable params: 0
Total mult-adds (G): 13.28
==========================================================================================
Input size (MB): 0.14
Forward/backward pass size (MB): 71.21
Params size (MB): 129.43
Estimated Total Size (MB): 200.78
==========================================================================================
文章出处登录后可见！
已经登录？立即刷新
