pytorch输出网络每层的参数和维度信息

本文主要是介绍pytorch输出网络每层的参数和维度信息，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

导读

tensorflow提供了一个summary函数可以让我们很方便的查看到网络每层的输出的shape信息和参数信息，pytorch中也提供了类似的方法

pytorch输出网络每层信息

安装torchsummary

pip install torchsummary

输出网络信息

summary函数介绍
model：网络模型
input_size：网络输入图片的shape，这里不用加batch_size进去
batch_size：batch_size参数，默认是-1
device：在GPU还是CPU上运行，默认是cuda在GPU上运行，如果想在CPU上执行将参数改为CPU即可

import torch.nn as nn
from torchsummary import summary#定义网络结构
net = nn.Sequential(nn.Conv2d(1,8,kernel_size=7),nn.MaxPool2d(2,stride=2),nn.ReLU(True),nn.Conv2d(8,10,kernel_size=5),nn.MaxPool2d(2,stride=2),nn.ReLU(True))#输出每层网络参数信息
summary(net,(1,28,28),batch_size=1,device="cpu")----------------------------------------------------------------Layer (type)               Output Shape         Param #
================================================================Conv2d-1             [1, 8, 22, 22]             400MaxPool2d-2             [1, 8, 11, 11]               0ReLU-3             [1, 8, 11, 11]               0Conv2d-4              [1, 10, 7, 7]           2,010MaxPool2d-5              [1, 10, 3, 3]               0ReLU-6              [1, 10, 3, 3]               0
================================================================
Total params: 2,410
Trainable params: 2,410
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.05
Params size (MB): 0.01
Estimated Total Size (MB): 0.06
----------------------------------------------------------------