首先,安装torchsummary
pip isntall torchsummary
下面是demo代码,其中(3, 100, 100)
分别为单张图片的通道数、高、宽
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)
输出结果如下,13个Conv2d(卷积层),3个Linear(全连接层),一共16层,vgg16没毛病
----------------------------------------------------------------
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
----------------------------------------------------------------