pytorch中获取模型input/output shape实例
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实例就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持猪先飞。
相关文章
pytorch nn.Conv2d()中的padding以及输出大小方式
今天小编就为大家分享一篇pytorch nn.Conv2d()中的padding以及输出大小方式,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧...2020-04-27Linux安装Pytorch1.8GPU(CUDA11.1)的实现
这篇文章主要介绍了Linux安装Pytorch1.8GPU(CUDA11.1)的实现,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友们下面随着小编来一起学习学习吧...2021-03-25- 这篇文章主要介绍了PyTorch一小时掌握之迁移学习篇,本文给大家介绍的非常详细,对大家的学习或工作具有一定的参考借鉴价值,需要的朋友可以参考下...2021-09-08
- 今天小编就为大家分享一篇pytorch 自定义卷积核进行卷积操作方式,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧...2020-05-06
- 这篇文章主要介绍了Pytorch之扩充tensor的操作,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧...2021-03-05
- 这篇文章主要介绍了解决pytorch 交叉熵损失输出为负数的问题,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧...2020-07-08
- 这篇文章主要介绍了JS双击变input框批量修改内容的实现代码,非常不错,具有参考借鉴价值,需要的朋友可以参考下...2017-01-09
- 这篇文章主要介绍了vue:el-input输入时限制输入的类型操作,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧...2020-08-05
pytorch绘制并显示loss曲线和acc曲线,LeNet5识别图像准确率
今天小编就为大家分享一篇pytorch绘制并显示loss曲线和acc曲线,LeNet5识别图像准确率,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧...2020-05-02- 这篇文章主要介绍了pytorch 实现冻结部分参数训练另一部分,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧...2021-03-27
vue element table中自定义一些input的验证操作
这篇文章主要介绍了vue element table中自定义一些input的验证操作,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧...2020-07-18从Pytorch模型pth文件中读取参数成numpy矩阵的操作
这篇文章主要介绍了从Pytorch模型pth文件中读取参数成numpy矩阵的操作,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧...2021-03-04Pytorch 的损失函数Loss function使用详解
今天小编就为大家分享一篇Pytorch 的损失函数Loss function使用详解,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧...2020-05-02如何清除IE10+ input X 文本框的叉叉和密码输入框的眼睛图标
从IE 10开始,type=”text” 的 input 在用户输入内容后,会自动产生一个小叉叉(X),方便用户点击清除已经输入的文本,下面通过本文给大家介绍下如何清除IE10+ input X 文本框的叉叉和密码输入框的眼睛图标...2017-01-09- 今天小编就为大家分享一篇pytorch中的上采样以及各种反操作,求逆操作详解,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧...2020-04-30
- 这篇文章主要介绍了pytorch中的squeeze函数、cat函数使用,具有很好的参考价值,希望对大家有所帮助。如有错误或未考虑完全的地方,望不吝赐教...2021-05-20
- 今天小编就为大家分享一篇Pytorch实现LSTM和GRU示例,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧...2020-04-27
- 这篇文章主要介绍了基于Pytorch版yolov5的滑块验证码破解思路详解,本文给大家介绍的非常详细,对大家的学习或工作具有一定的参考借鉴价值,需要的朋友可以参考下...2021-02-25
解决Pytorch dataloader时报错每个tensor维度不一样的问题
这篇文章主要介绍了解决Pytorch dataloader时报错每个tensor维度不一样的问题,具有很好的参考价值,希望对大家有所帮助。如有错误或未考虑完全的地方,望不吝赐教...2021-05-28pytorch中交叉熵损失(nn.CrossEntropyLoss())的计算过程详解
今天小编就为大家分享一篇pytorch中交叉熵损失(nn.CrossEntropyLoss())的计算过程详解,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧...2020-05-02