pytorch中獲取模型input/output shape實例
更新時間：2019年12月30日 09:13:24 作者：mylibrary1
今天小編就為大家分享一篇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實例就是小編分享給大家的全部內容了，希望能給大家一個參考，也希望大家多多支持腳本之家。
相關文章
史上十大最經典小霸王游戲中魂斗羅只能排在第二，那么第一是誰？最經典最風靡的當屬超級瑪麗，那個戴帽子的大胡子穿著背帶褲的馬里奧哪個不認得，小編帶你用python實現超級瑪麗緬懷童年2021-09-09
在本篇內容中我們給大家整理了關于python面向對象編程基礎的觀點分析以及知識點整理，有需要的朋友們學習下。2019-05-05
今天分享個如何簡單處理滑動圖片的驗證碼的案例，主要是使用Python爬蟲破解滑動驗證碼的相關實現代碼，感興趣的朋友跟隨小編一起看看吧2021-05-05
Python使用flask-caching緩存數據的示例代碼
Flask-Caching 是 Flask 的一個擴展,為任何 Flask 應用程序添加了對各種后端的緩存支持,它基于 cachelib 運行,并通過統(tǒng)一的 API 支持 werkzeug 的所有原始緩存后端,本文給大家介紹了Python使用flask-caching緩存數據,需要的朋友可以參考下2024-12-12
詳解Django-restframework 之頻率源碼分析
這篇文章主要介紹了Django-restframework 之頻率源碼分析，小編覺得挺不錯的，現在分享給大家，也給大家做個參考。一起跟隨小編過來看看吧2019-02-02
教你用Type Hint提高Python程序開發(fā)效率
本文通過介紹和實例教大家如何利用Type Hint來提升Python程序開發(fā)效率，對大家使用python開發(fā)很有幫助，有需要的參考學習。2016-08-08
這篇文章主要介紹了Pycharm如何導入python文件及解決報錯問題,本文通過示例截圖相結合給大家介紹的非常詳細，對大家的學習或工作具有一定的參考借鑒價值，需要的朋友可以參考下2020-05-05
這篇文章主要介紹了Python內置數據類型list各方法的性能測試過程解析,文中通過示例代碼介紹的非常詳細，對大家的學習或者工作具有一定的參考學習價值,需要的朋友可以參考下2020-01-01
這篇文章主要介紹了Python flask框架實現瀏覽器點擊自定義跳轉頁面,文中通過示例代碼介紹的非常詳細，對大家的學習或者工作具有一定的參考學習價值,需要的朋友可以參考下2020-06-06
typing.Dict和Dict的區(qū)別及它們在Python中的用途小結
當在 Python 函數中聲明一個 dictionary 作為參數時，我們一般會把 key 和 value 的數據類型聲明為全局變量，而不是局部變量。，這篇文章主要介紹了typing.Dict和Dict的區(qū)別及它們在Python中的用途小結,需要的朋友可以參考下2023-06-06