Generating Keras-like model summary in PyTorch
Published in
6 min readOct 25, 2019
If you are an ardent Keras user and are recently moving to PyTorch, I am pretty sure you would be missing so many awesome features of keras. Salute to Francois Chollet for Keras.
One such amazing feature is quickly getting the model summary and making the sense of your developed architecture at once using the Keras model.summary()
from keras.datasets import mnist
from keras.models import Sequential
from keras.layers import Dense, Dropout, Flatten
from keras.layers import Conv2D, MaxPooling2Dmodel = Sequential()
model.add(Conv2D(32, kernel_size=(3, 3),
activation='relu',
input_shape=input_shape))
model.add(Conv2D(64, (3, 3), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Dropout(0.25))
model.add(Flatten())
model.add(Dense(128, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(num_classes,activation='softmax'))model.summary()_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
conv2d_4 (Conv2D) (None, 222, 222, 32) 896
_________________________________________________________________
conv2d_5 (Conv2D) (None, 220, 220, 64) 18496
_________________________________________________________________
max_pooling2d_1 (MaxPooling2 (None, 110, 110, 64) 0
_________________________________________________________________
dropout_1 (Dropout) (None, 110, 110, 64) 0
_________________________________________________________________
flatten_1 (Flatten) (None, 774400) 0
_________________________________________________________________
dense_10 (Dense) (None, 128) 99123328
_________________________________________________________________
dropout_2 (Dropout) (None, 128) 0
_________________________________________________________________
dense_11 (Dense) (None, 10) 1290
=================================================================
Total params: 99,144,010
Trainable params: 99,144,010
Non-trainable params: 0
_________________________________________________________________
The above snippet is from Keras and we can see how easily we can see the entire model summary with output shape and the number of parameters.
Is there some similar function in PyTorch??? Yess!!!
It’s the torchsummary library
Installation
Installation is pretty straightforward
pip install torchsummaryUsing torchsummary
summary(model,input_shape)Let me show through an example and for this elucidation I will import a pretrained Alexnet model trained on ImageNet dataset
from torchvision import models
from torchsummary import summary# import the pretrained alexnet
model = models.alexnet(pretrained=True)
summary(model,(3,224,224))# And you should see something like this
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 64, 55, 55] 23,296
ReLU-2 [-1, 64, 55, 55] 0
MaxPool2d-3 [-1, 64, 27, 27] 0
Conv2d-4 [-1, 192, 27, 27] 307,392
ReLU-5 [-1, 192, 27, 27] 0
MaxPool2d-6 [-1, 192, 13, 13] 0
Conv2d-7 [-1, 384, 13, 13] 663,936
ReLU-8 [-1, 384, 13, 13] 0
Conv2d-9 [-1, 256, 13, 13] 884,992
ReLU-10 [-1, 256, 13, 13] 0
Conv2d-11 [-1, 256, 13, 13] 590,080
ReLU-12 [-1, 256, 13, 13] 0
MaxPool2d-13 [-1, 256, 6, 6] 0
Dropout-14 [-1, 9216] 0
Linear-15 [-1, 4096] 37,752,832
ReLU-16 [-1, 4096] 0
Dropout-17 [-1, 4096] 0
Linear-18 [-1, 4096] 16,781,312
ReLU-19 [-1, 4096] 0
Linear-20 [-1, 1000] 4,097,000
================================================================
Total params: 61,100,840
Trainable params: 61,100,840
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.57
Forward/backward pass size (MB): 8.31
Params size (MB): 233.08
Estimated Total Size (MB): 241.96
----------------------------------------------------------------
with Cuda
If you have cuda you may need to export your model to cuda if you get a run time error as shown below
model = models.alexnet(pretrained=True)
summary(model,(3,224,224))RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same#instead do thismodel = models.alexnet(pretrained=True)
model.cuda()
summary(model,(3,224,224))----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 64, 55, 55] 23,296
ReLU-2 [-1, 64, 55, 55] 0
MaxPool2d-3 [-1, 64, 27, 27] 0
Conv2d-4 [-1, 192, 27, 27] 307,392
ReLU-5 [-1, 192, 27, 27] 0
MaxPool2d-6 [-1, 192, 13, 13] 0
Conv2d-7 [-1, 384, 13, 13] 663,936
ReLU-8 [-1, 384, 13, 13] 0
Conv2d-9 [-1, 256, 13, 13] 884,992
ReLU-10 [-1, 256, 13, 13] 0
Conv2d-11 [-1, 256, 13, 13] 590,080
ReLU-12 [-1, 256, 13, 13] 0
MaxPool2d-13 [-1, 256, 6, 6] 0
AdaptiveAvgPool2d-14 [-1, 256, 6, 6] 0
Dropout-15 [-1, 9216] 0
Linear-16 [-1, 4096] 37,752,832
ReLU-17 [-1, 4096] 0
Dropout-18 [-1, 4096] 0
Linear-19 [-1, 4096] 16,781,312
ReLU-20 [-1, 4096] 0
Linear-21 [-1, 1000] 4,097,000
================================================================
Total params: 61,100,840
Trainable params: 61,100,840
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.57
Forward/backward pass size (MB): 8.38
Params size (MB): 233.08
Estimated Total Size (MB): 242.03
----------------------------------------------------------------
Summary of a 3DCNN Model
Here the input shape will have an additional dimension of number of frames. I am using 3D CNN ResNet for video classification from Kensho Hara-https://github.com/kenshohara/video-classification-3d-cnn-pytorch
It is a 3D ResNet-18 architecture with input shape of (3,16,112,112)
model = generate_model(opt)
summary(model,(3,16,112,112))----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv3d-1 [-1, 64, 16, 56, 56] 65,856
BatchNorm3d-2 [-1, 64, 16, 56, 56] 128
ReLU-3 [-1, 64, 16, 56, 56] 0
MaxPool3d-4 [-1, 64, 8, 28, 28] 0
Conv3d-5 [-1, 64, 8, 28, 28] 110,592
BatchNorm3d-6 [-1, 64, 8, 28, 28] 128
ReLU-7 [-1, 64, 8, 28, 28] 0
Conv3d-8 [-1, 64, 8, 28, 28] 110,592
BatchNorm3d-9 [-1, 64, 8, 28, 28] 128
ReLU-10 [-1, 64, 8, 28, 28] 0
BasicBlock-11 [-1, 64, 8, 28, 28] 0
Conv3d-12 [-1, 64, 8, 28, 28] 110,592
BatchNorm3d-13 [-1, 64, 8, 28, 28] 128
ReLU-14 [-1, 64, 8, 28, 28] 0
Conv3d-15 [-1, 64, 8, 28, 28] 110,592
BatchNorm3d-16 [-1, 64, 8, 28, 28] 128
ReLU-17 [-1, 64, 8, 28, 28] 0
BasicBlock-18 [-1, 64, 8, 28, 28] 0
Conv3d-19 [-1, 128, 4, 14, 14] 221,184
BatchNorm3d-20 [-1, 128, 4, 14, 14] 256
ReLU-21 [-1, 128, 4, 14, 14] 0
Conv3d-22 [-1, 128, 4, 14, 14] 442,368
BatchNorm3d-23 [-1, 128, 4, 14, 14] 256
ReLU-24 [-1, 128, 4, 14, 14] 0
BasicBlock-25 [-1, 128, 4, 14, 14] 0
Conv3d-26 [-1, 128, 4, 14, 14] 442,368
BatchNorm3d-27 [-1, 128, 4, 14, 14] 256
ReLU-28 [-1, 128, 4, 14, 14] 0
Conv3d-29 [-1, 128, 4, 14, 14] 442,368
BatchNorm3d-30 [-1, 128, 4, 14, 14] 256
ReLU-31 [-1, 128, 4, 14, 14] 0
BasicBlock-32 [-1, 128, 4, 14, 14] 0
Conv3d-33 [-1, 256, 2, 7, 7] 884,736
BatchNorm3d-34 [-1, 256, 2, 7, 7] 512
ReLU-35 [-1, 256, 2, 7, 7] 0
Conv3d-36 [-1, 256, 2, 7, 7] 1,769,472
BatchNorm3d-37 [-1, 256, 2, 7, 7] 512
ReLU-38 [-1, 256, 2, 7, 7] 0
BasicBlock-39 [-1, 256, 2, 7, 7] 0
Conv3d-40 [-1, 256, 2, 7, 7] 1,769,472
BatchNorm3d-41 [-1, 256, 2, 7, 7] 512
ReLU-42 [-1, 256, 2, 7, 7] 0
Conv3d-43 [-1, 256, 2, 7, 7] 1,769,472
BatchNorm3d-44 [-1, 256, 2, 7, 7] 512
ReLU-45 [-1, 256, 2, 7, 7] 0
BasicBlock-46 [-1, 256, 2, 7, 7] 0
Conv3d-47 [-1, 512, 1, 4, 4] 3,538,944
BatchNorm3d-48 [-1, 512, 1, 4, 4] 1,024
ReLU-49 [-1, 512, 1, 4, 4] 0
Conv3d-50 [-1, 512, 1, 4, 4] 7,077,888
BatchNorm3d-51 [-1, 512, 1, 4, 4] 1,024
ReLU-52 [-1, 512, 1, 4, 4] 0
BasicBlock-53 [-1, 512, 1, 4, 4] 0
Conv3d-54 [-1, 512, 1, 4, 4] 7,077,888
BatchNorm3d-55 [-1, 512, 1, 4, 4] 1,024
ReLU-56 [-1, 512, 1, 4, 4] 0
Conv3d-57 [-1, 512, 1, 4, 4] 7,077,888
BatchNorm3d-58 [-1, 512, 1, 4, 4] 1,024
ReLU-59 [-1, 512, 1, 4, 4] 0
BasicBlock-60 [-1, 512, 1, 4, 4] 0
AvgPool3d-61 [-1, 512, 1, 1, 1] 0
Linear-62 [-1, 400] 205,200
ResNet-63 [-1, 400] 0
================================================================
Total params: 33,235,280
Trainable params: 33,235,280
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 2.30
Forward/backward pass size (MB): 133.72
Params size (MB): 126.78
Estimated Total Size (MB): 262.80
----------------------------------------------------------------
I hope this post was useful!!
Till then, Keep Learning! Keep Exploring Neurons!
If you find my articles helpful and wish to support them — Buy me a Coffee
