pytorch中获取模型input/outputshape实例-创新互联
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)
另外有需要云服务器可以了解下创新互联scvps.cn,海内外云服务器15元起步,三天无理由+7*72小时售后在线,公司持有idc许可证,提供“云服务器、裸金属服务器、高防服务器、香港服务器、美国服务器、虚拟主机、免备案服务器”等云主机租用服务以及企业上云的综合解决方案,具有“安全稳定、简单易用、服务可用性高、性价比高”等特点与优势,专为企业上云打造定制,能够满足用户丰富、多元化的应用场景需求。
网站名称:pytorch中获取模型input/outputshape实例-创新互联
分享路径:http://myzitong.com/article/djeihh.html