pytorch 小技巧2

1、网络参数显示——summary

#官网：https://github.com/sksq96/pytorch-summary
from torchsummary import summary
summary(your_model, input_size=(channels, H, W))

用法：

from torchsummary import  summary

#中间是一堆网络的结构

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # PyTorch v0.4.0
## 输出我们的网络结构
MyConvnet = ConvNet().to(device)
summary(MyConvnet,(1,28,28))

显示的情况：

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 16, 28, 28]             160
              ReLU-2           [-1, 16, 28, 28]               0
         AvgPool2d-3           [-1, 16, 14, 14]               0
            Conv2d-4           [-1, 32, 14, 14]           4,640
              ReLU-5           [-1, 32, 14, 14]               0
         MaxPool2d-6             [-1, 32, 7, 7]               0
            Linear-7                  [-1, 128]         200,832
              ReLU-8                  [-1, 128]               0
            Linear-9                   [-1, 64]           8,256
             ReLU-10                   [-1, 64]               0
           Linear-11                   [-1, 10]             650
================================================================
Total params: 214,538
Trainable params: 214,538
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.33
Params size (MB): 0.82
Estimated Total Size (MB): 1.15
----------------------------------------------------------------