【Pytorch】Quickly output network model information method

First, installtorchsummary

pip isntall torchsummary

The following is the demo code, which (3, 100, 100)is the number of channels, height, and width of a single picture.

from torchsummary import summary
import torchvision
import torch

def model_info(model):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    backbone = model.to(device)
    summary(backbone, (3, 100, 100))

model = torchvision.models.vgg16(pretrained=False)
model_info(model)

The output result is as follows, 13 Conv2d (convolutional layer), 3 Linear (fully connected layer), a total of 16 layers, vgg16 is fine

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 64, 100, 100]           1,792
              ReLU-2         [-1, 64, 100, 100]               0
            Conv2d-3         [-1, 64, 100, 100]          36,928
              ReLU-4         [-1, 64, 100, 100]               0
         MaxPool2d-5           [-1, 64, 50, 50]               0
            Conv2d-6          [-1, 128, 50, 50]          73,856
              ReLU-7          [-1, 128, 50, 50]               0
            Conv2d-8          [-1, 128, 50, 50]         147,584
              ReLU-9          [-1, 128, 50, 50]               0
        MaxPool2d-10          [-1, 128, 25, 25]               0
           Conv2d-11          [-1, 256, 25, 25]         295,168
             ReLU-12          [-1, 256, 25, 25]               0
           Conv2d-13          [-1, 256, 25, 25]         590,080
             ReLU-14          [-1, 256, 25, 25]               0
           Conv2d-15          [-1, 256, 25, 25]         590,080
             ReLU-16          [-1, 256, 25, 25]               0
        MaxPool2d-17          [-1, 256, 12, 12]               0
           Conv2d-18          [-1, 512, 12, 12]       1,180,160
             ReLU-19          [-1, 512, 12, 12]               0
           Conv2d-20          [-1, 512, 12, 12]       2,359,808
             ReLU-21          [-1, 512, 12, 12]               0
           Conv2d-22          [-1, 512, 12, 12]       2,359,808
             ReLU-23          [-1, 512, 12, 12]               0
        MaxPool2d-24            [-1, 512, 6, 6]               0
           Conv2d-25            [-1, 512, 6, 6]       2,359,808
             ReLU-26            [-1, 512, 6, 6]               0
           Conv2d-27            [-1, 512, 6, 6]       2,359,808
             ReLU-28            [-1, 512, 6, 6]               0
           Conv2d-29            [-1, 512, 6, 6]       2,359,808
             ReLU-30            [-1, 512, 6, 6]               0
        MaxPool2d-31            [-1, 512, 3, 3]               0
AdaptiveAvgPool2d-32            [-1, 512, 7, 7]               0
           Linear-33                 [-1, 4096]     102,764,544
             ReLU-34                 [-1, 4096]               0
          Dropout-35                 [-1, 4096]               0
           Linear-36                 [-1, 4096]      16,781,312
             ReLU-37                 [-1, 4096]               0
          Dropout-38                 [-1, 4096]               0
           Linear-39                 [-1, 1000]       4,097,000
================================================================
Total params: 138,357,544
Trainable params: 138,357,544
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.11
Forward/backward pass size (MB): 43.51
Params size (MB): 527.79
Estimated Total Size (MB): 571.42
----------------------------------------------------------------