After using PyTorch deep learning to build a model, if you want to view the model structure, you can directly use the print(model) function to print it. However, the output result is not particularly intuitive. After checking, I found that there is a model visualization tool that can output model.summary() similar to keras style. Record it here for later reference.
PyTorch print model structure, output dimensions and parameter information (torchsummary)
Install torchsummary
pip install torchsummary
output network information
summary function introduction
model: Network model
input_size: The shape of the network input image, there is no need to add batch_size here
batch_size: batch_size parameter, the default is -1
device: run on GPU or CPU, the default is cuda running on GPU, if you want to execute on CPU, change the parameter to CPU can
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()
output:
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’
Adding the old summary to LSTM and the like will report an error, and a new summary is required
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
==========================================================================================