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

以上這篇pytorch中獲取模型input/output shape例項就是小編分享給大家的全部內容了，希望能給大家一個參考，也希望大家多多支援我們。