model summary

use torchsummary package to get summary of the model to see the output shape and the number of parameters in each layer.

 

1. install the torchsummary package:

pip install torchsummary

2. let's get the model summary using torchsummary:

from torchsummary import summary
summary(model, input_size=(channels, H, W))

 

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

 

※ reference: github.com/sksq96/pytorch-summary