使用torchsummary打印出神经网络的形状和参数大小

首先安装torchsummary

pip install torchsummary

直接另存为py文件

from torchsummary import summary
import torchvision.models as models
model = models.vgg16_bn()
model = model.cuda()
summary(model, input_size=(3,224,224), batch_size=-1, device='cuda')

结果如下