我们知道,Keras有一个非常有好的功能是summary,可以打印显示网络结构和参数,一目了然。但是,Pytorch本身好像不支持这一点。不过,幸好有一个工具叫torchsummary,可以实现和Keras几乎一样的效果。
pip install torchsummary
然后我们定义好网络结构之后,就可以用summary来打印显示了。假设我们定义的网络结构是一个叫Generator的类。
import torch from torchsummary import summary # 需要使用device来指定网络在GPU还是CPU运行
device= torch.device('cuda' if torch.cuda.is_available() else 'cpu') netG_A2B =
Generator(3, 3).to(device) summary(netG_A2B, input_size=(3, 256, 256))
之后,就可以打印网络结构了。一个示例结构如下:
---------------------------------------------------------------- Layer (type)
Output Shape Param# ============================================================
==== Conv2d-1 [-1, 64, 224, 224] 1,792 ReLU-2 [-1, 64, 224, 224] 0 Conv2d-3 [-1,
64, 224, 224] 36,928 ReLU-4 [-1, 64, 224, 224] 0 MaxPool2d-5 [-1, 64, 112, 112]
0 Conv2d-6 [-1, 128, 112, 112] 73,856 ReLU-7 [-1, 128, 112, 112] 0 Conv2d-8 [-1,
128, 112, 112] 147,584 ReLU-9 [-1, 128, 112, 112] 0 MaxPool2d-10 [-1, 128, 56,
56] 0 Conv2d-11 [-1, 256, 56, 56] 295,168 ReLU-12 [-1, 256, 56, 56] 0 Conv2d-13
[-1, 256, 56, 56] 590,080 ReLU-14 [-1, 256, 56, 56] 0 Conv2d-15 [-1, 256, 56, 56
] 590,080 ReLU-16 [-1, 256, 56, 56] 0 MaxPool2d-17 [-1, 256, 28, 28] 0 Conv2d-18
[-1, 512, 28, 28] 1,180,160 ReLU-19 [-1, 512, 28, 28] 0 Conv2d-20 [-1, 512, 28,
28] 2,359,808 ReLU-21 [-1, 512, 28, 28] 0 Conv2d-22 [-1, 512, 28, 28] 2,359,808
ReLU-23 [-1, 512, 28, 28] 0 MaxPool2d-24 [-1, 512, 14, 14] 0 Conv2d-25 [-1, 512,
14, 14] 2,359,808 ReLU-26 [-1, 512, 14, 14] 0 Conv2d-27 [-1, 512, 14, 14] 2,359,
808 ReLU-28 [-1, 512, 14, 14] 0 Conv2d-29 [-1, 512, 14, 14] 2,359,808 ReLU-30 [-
1, 512, 14, 14] 0 MaxPool2d-31 [-1, 512, 7, 7] 0 Linear-32 [-1, 4096] 102,764,
544 ReLU-33 [-1, 4096] 0 Dropout-34 [-1, 4096] 0 Linear-35 [-1, 4096] 16,781,312
ReLU-36 [-1, 4096] 0 Dropout-37 [-1, 4096] 0 Linear-38 [-1, 1000] 4,097,000 ==
============================================================== Total params: 138
,357,544 Trainable params: 138,357,544 Non-trainable params: 0 -----------------
----------------------------------------------- Input size (MB): 0.57 Forward/
backwardpass size (MB): 218.59 Params size (MB): 527.79 Estimated Total Size (MB
): 746.96 ----------------------------------------------------------------

技术
©2019-2020 Toolsou All rights reserved,
车主无忧:为什么放弃开源Kafka?必传之作!Alibaba内部出品Redis深度笔记及源码宝典统信UOS首次公布软件适配:QQ、微信、迅雷都有了阿里开发10年技术核心总结,Springboot+Redis文档,送给努力上进的程序员python画爱心Docker容器数据卷详解(共享数据)为何华为的5G专利高居第一名,却还被高通要求缴纳专利费?Spring循环依赖三级缓存是否可以减少为二级缓存?C#中字典的排序方法面试还不会Spring?阿里P8总结的100道面试解析,让你实锤面试官