pytorch中获取模型input/output shape实例


Posted in Python onDecember 30, 2019

Pytorch官方目前无法像tensorflow, caffe那样直接给出shape信息,详见

https://github.com/pytorch/pytorch/pull/3043

以下代码算一种workaround。由于CNN, RNN等模块实现不一样,添加其他模块支持可能需要改代码。

例如RNN中bias是bool类型,其权重也不是存于weight属性中,不过我们只关注shape够用了。

该方法必须构造一个输入调用forward后(model(x)调用)才可获取shape

#coding:utf-8
from collections import OrderedDict
import torch
from torch.autograd import Variable
import torch.nn as nn
import models.crnn as crnn
import json
 
 
def get_output_size(summary_dict, output):
 if isinstance(output, tuple):
 for i in xrange(len(output)):
  summary_dict[i] = OrderedDict()
  summary_dict[i] = get_output_size(summary_dict[i],output[i])
 else:
 summary_dict['output_shape'] = list(output.size())
 return summary_dict
 
def summary(input_size, model):
 def register_hook(module):
 def hook(module, input, output):
  class_name = str(module.__class__).split('.')[-1].split("'")[0]
  module_idx = len(summary)
 
  m_key = '%s-%i' % (class_name, module_idx+1)
  summary[m_key] = OrderedDict()
  summary[m_key]['input_shape'] = list(input[0].size())
  summary[m_key] = get_output_size(summary[m_key], output)
 
  params = 0
  if hasattr(module, 'weight'):
  params += torch.prod(torch.LongTensor(list(module.weight.size())))
  if module.weight.requires_grad:
   summary[m_key]['trainable'] = True
  else:
   summary[m_key]['trainable'] = False
  #if hasattr(module, 'bias'):
  # params += torch.prod(torch.LongTensor(list(module.bias.size())))
 
  summary[m_key]['nb_params'] = params
  
 if not isinstance(module, nn.Sequential) and \
  not isinstance(module, nn.ModuleList) and \
  not (module == model):
  hooks.append(module.register_forward_hook(hook))
 
 # check if there are multiple inputs to the network
 if isinstance(input_size[0], (list, tuple)):
 x = [Variable(torch.rand(1,*in_size)) for in_size in input_size]
 else:
 x = Variable(torch.rand(1,*input_size))
 
 # create properties
 summary = OrderedDict()
 hooks = []
 # register hook
 model.apply(register_hook)
 # make a forward pass
 model(x)
 # remove these hooks
 for h in hooks:
 h.remove()
 
 return summary
 
crnn = crnn.CRNN(32, 1, 3755, 256, 1)
x = summary([1,32,128],crnn)
print json.dumps(x)

以pytorch版CRNN为例,输出shape如下

{
"Conv2d-1": {
"input_shape": [1, 1, 32, 128],
"output_shape": [1, 64, 32, 128],
"trainable": true,
"nb_params": 576
},
"ReLU-2": {
"input_shape": [1, 64, 32, 128],
"output_shape": [1, 64, 32, 128],
"nb_params": 0
},
"MaxPool2d-3": {
"input_shape": [1, 64, 32, 128],
"output_shape": [1, 64, 16, 64],
"nb_params": 0
},
"Conv2d-4": {
"input_shape": [1, 64, 16, 64],
"output_shape": [1, 128, 16, 64],
"trainable": true,
"nb_params": 73728
},
"ReLU-5": {
"input_shape": [1, 128, 16, 64],
"output_shape": [1, 128, 16, 64],
"nb_params": 0
},
"MaxPool2d-6": {
"input_shape": [1, 128, 16, 64],
"output_shape": [1, 128, 8, 32],
"nb_params": 0
},
"Conv2d-7": {
"input_shape": [1, 128, 8, 32],
"output_shape": [1, 256, 8, 32],
"trainable": true,
"nb_params": 294912
},
"BatchNorm2d-8": {
"input_shape": [1, 256, 8, 32],
"output_shape": [1, 256, 8, 32],
"trainable": true,
"nb_params": 256
},
"ReLU-9": {
"input_shape": [1, 256, 8, 32],
"output_shape": [1, 256, 8, 32],
"nb_params": 0
},
"Conv2d-10": {
"input_shape": [1, 256, 8, 32],
"output_shape": [1, 256, 8, 32],
"trainable": true,
"nb_params": 589824
},
"ReLU-11": {
"input_shape": [1, 256, 8, 32],
"output_shape": [1, 256, 8, 32],
"nb_params": 0
},
"MaxPool2d-12": {
"input_shape": [1, 256, 8, 32],
"output_shape": [1, 256, 4, 33],
"nb_params": 0
},
"Conv2d-13": {
"input_shape": [1, 256, 4, 33],
"output_shape": [1, 512, 4, 33],
"trainable": true,
"nb_params": 1179648
},
"BatchNorm2d-14": {
"input_shape": [1, 512, 4, 33],
"output_shape": [1, 512, 4, 33],
"trainable": true,
"nb_params": 512
},
"ReLU-15": {
"input_shape": [1, 512, 4, 33],
"output_shape": [1, 512, 4, 33],
"nb_params": 0
},
"Conv2d-16": {
"input_shape": [1, 512, 4, 33],
"output_shape": [1, 512, 4, 33],
"trainable": true,
"nb_params": 2359296
},
"ReLU-17": {
"input_shape": [1, 512, 4, 33],
"output_shape": [1, 512, 4, 33],
"nb_params": 0
},
"MaxPool2d-18": {
"input_shape": [1, 512, 4, 33],
"output_shape": [1, 512, 2, 34],
"nb_params": 0
},
"Conv2d-19": {
"input_shape": [1, 512, 2, 34],
"output_shape": [1, 512, 1, 33],
"trainable": true,
"nb_params": 1048576
},
"BatchNorm2d-20": {
"input_shape": [1, 512, 1, 33],
"output_shape": [1, 512, 1, 33],
"trainable": true,
"nb_params": 512
},
"ReLU-21": {
"input_shape": [1, 512, 1, 33],
"output_shape": [1, 512, 1, 33],
"nb_params": 0
},
"LSTM-22": {
"input_shape": [33, 1, 512],
"0": {
"output_shape": [33, 1, 512]
},
"1": {
"0": {
"output_shape": [2, 1, 256]
},
"1": {
"output_shape": [2, 1, 256]
}
},
"nb_params": 0
},
"Linear-23": {
"input_shape": [33, 512],
"output_shape": [33, 256],
"trainable": true,
"nb_params": 131072
},
"BidirectionalLSTM-24": {
"input_shape": [33, 1, 512],
"output_shape": [33, 1, 256],
"nb_params": 0
},
"LSTM-25": {
"input_shape": [33, 1, 256],
"0": {
"output_shape": [33, 1, 512]
},
"1": {
"0": {
"output_shape": [2, 1, 256]
},
"1": {
"output_shape": [2, 1, 256]
}
},
"nb_params": 0
},
"Linear-26": {
"input_shape": [33, 512],
"output_shape": [33, 3755],
"trainable": true,
"nb_params": 1922560
},
"BidirectionalLSTM-27": {
"input_shape": [33, 1, 256],
"output_shape": [33, 1, 3755],
"nb_params": 0
}
}

以上这篇pytorch中获取模型input/output shape实例就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持三水点靠木。

Python 相关文章推荐
基于python编写的微博应用
Oct 17 Python
使用Python编写vim插件的简单示例
Apr 17 Python
详解字典树Trie结构及其Python代码实现
Jun 03 Python
python2 与 python3 实现共存的方法
Jul 12 Python
JavaScript中的模拟事件和自定义事件实例分析
Jul 27 Python
python实现静态web服务器
Sep 03 Python
Python 类的魔法属性用法实例分析
Nov 21 Python
运行tensorflow python程序,限制对GPU和CPU的占用操作
Feb 06 Python
详解python polyscope库的安装和例程
Nov 13 Python
Python如何导出导入所有依赖包详解
Jun 08 Python
Python数据处理的三个实用技巧分享
Apr 01 Python
Python如何加载模型并查看网络
Jul 15 Python
Python读取csv文件实例解析
Dec 30 #Python
Pytorch Tensor的统计属性实例讲解
Dec 30 #Python
PyTorch中permute的用法详解
Dec 30 #Python
python实现多进程按序号批量修改文件名的方法示例
Dec 30 #Python
Pytorch Tensor基本数学运算详解
Dec 30 #Python
python垃圾回收机制(GC)原理解析
Dec 30 #Python
利用Python代码实现一键抠背景功能
Dec 29 #Python
You might like
第十节 抽象方法和抽象类 [10]
2006/10/09 PHP
说明的比较细的php 正则学习实例
2008/07/30 PHP
关于php操作mysql执行数据库查询的一些常用操作汇总
2013/06/24 PHP
PHP的password_hash()使用实例
2014/03/17 PHP
php出现内存位置访问无效错误问题解决方法
2014/08/16 PHP
常见的四种POST 提交数据方式(小总结)
2015/10/08 PHP
分析PHP中单双引号的误区和双引号小隐患
2016/07/19 PHP
php遍历、读取文件夹中图片并分页显示图片的方法
2016/11/15 PHP
PHP 文件上传后端处理实用技巧方法
2017/01/06 PHP
laravel实现图片上传预览,及编辑时可更换图片,并实时变化的例子
2019/11/14 PHP
某页码显示的helper 少量调整,另附js版
2010/09/12 Javascript
js登录弹出层特效
2014/03/07 Javascript
js确认删除对话框适用于a标签及submit
2014/07/10 Javascript
Nodejs初级阶段之express
2015/11/23 NodeJs
你一定会收藏的Nodejs代码片段
2016/02/04 NodeJs
jQuery实现手机自定义弹出输入框
2016/06/13 Javascript
NodeJs使用Mysql模块实现事务处理实例
2017/05/31 NodeJs
js实现日期显示的一些操作(实例讲解)
2017/07/27 Javascript
js禁止表单重复提交
2017/08/29 Javascript
深入浅析JavaScript中的RegExp对象
2017/09/18 Javascript
浅谈Node.js之异步流控制
2017/10/25 Javascript
vue.js根据代码运行环境选择baseurl的方法
2018/02/28 Javascript
使用vue-cli导入Element UI组件的方法
2018/05/16 Javascript
详解js中Array的方法及技巧
2018/09/12 Javascript
微信小程序下拉刷新PullDownRefresh的使用方法
2018/11/29 Javascript
小程序日历控件使用方法详解
2018/12/29 Javascript
30分钟精通React今年最劲爆的新特性——React Hooks
2019/03/11 Javascript
vue登录注册实例详解
2019/09/14 Javascript
解决antd 表单设置默认值initialValue后验证失效的问题
2020/11/02 Javascript
Python使用PyGreSQL操作PostgreSQL数据库教程
2014/07/30 Python
python求质数的3种方法
2018/09/28 Python
完美解决pycharm导入自己写的py文件爆红问题
2020/02/12 Python
CSS3利用text-shadow属性实现多种效果的文字样式展现方法
2016/08/25 HTML / CSS
WatchShop法国:英国排名第一的独立手表零售商
2020/02/17 全球购物
亮剑精神演讲稿
2014/05/23 职场文书