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许可证，提供“云服务器、裸金属服务器、高防服务器、香港服务器、美国服务器、虚拟主机、免备案服务器”等云主机租用服务以及企业上云的综合解决方案，具有“安全稳定、简单易用、服务可用性高、性价比高”等特点与优势，专为企业上云打造定制，能够满足用户丰富、多元化的应用场景需求。
