torchsummary工具的使用

安装：

pip install torchsummary

使用：输入为模型、输入尺寸、批数量、设备

from torchsummary import summary

summary(model, input_size, batch_size=-1, device="cuda")

使用时需要注意，默认device=‘cuda’，如果是在‘cpu’，那么就需要更改。不匹配就会出现下面的错误：

RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same

使用：

import torch
from torchvision import models, transforms
from torchsummary import summary

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# device = torch.device('cpu')
net = models.googlenet(pretrained=True).to(device)
summary(net, (3, 224, 224))

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 64, 112, 112]           9,408
       BatchNorm2d-2         [-1, 64, 112, 112]             128
       BasicConv2d-3         [-1, 64, 112, 112]               0
         MaxPool2d-4           [-1, 64, 56, 56]               0
            Conv2d-5           [-1, 64, 56, 56]           4,096
       BatchNorm2d-6           [-1, 64, 56, 56]             128
       BasicConv2d-7           [-1, 64, 56, 56]               0
            Conv2d-8          [-1, 192, 56, 56]         110,592
       BatchNorm2d-9          [-1, 192, 56, 56]             384
      BasicConv2d-10          [-1, 192, 56, 56]               0
        MaxPool2d-11          [-1, 192, 28, 28]               0
           Conv2d-12           [-1, 64, 28, 28]          12,288
      BatchNorm2d-13           [-1, 64, 28, 28]             128
      BasicConv2d-14           [-1, 64, 28, 28]               0
           Conv2d-15           [-1, 96, 28, 28]          18,432
      BatchNorm2d-16           [-1, 96, 28, 28]             192
      BasicConv2d-17           [-1, 96, 28, 28]               0
           Conv2d-18          [-1, 128, 28, 28]         110,592
      BatchNorm2d-19          [-1, 128, 28, 28]             256
      BasicConv2d-20          [-1, 128, 28, 28