tensorflow 模型權重導出實例
更新時間:2020年01月24日 09:34:01 作者:張中星
今天小編就為大家分享一篇tensorflow 模型權重導出實例,具有很好的參考價值,希望對大家有所幫助。一起跟隨小編過來看看吧
tensorflow在保存權重模型時多使用tf.train.Saver().save 函數進行權重保存,保存的ckpt文件無法直接打開,不利于將模型權重導入到其他框架使用(如Caffe、Keras等)。
好在tensorflow提供了相關函數 tf.train.NewCheckpointReader 可以對ckpt文件進行權重查看,因此可以通過該函數進行數據導出。
import tensorflow as tf
import h5py
cpktLogFileName = r'./checkpoint/checkpoint' #cpkt 文件路徑
with open(cpktLogFileName, 'r') as f:
#權重節(jié)點往往會保留多個epoch的數據,此處獲取最后的權重數據
cpktFileName = f.readline().split('"')[1]
h5FileName = r'./model/net_classification.h5'
reader = tf.train.NewCheckpointReader(cpktFileName)
f = h5py.File(h5FileName, 'w')
t_g = None
for key in sorted(reader.get_variable_to_shape_map()):
# 權重名稱需根據自己網絡名稱自行修改
if key.endswith('w') or key.endswith('biases'):
keySplits = key.split(r'/')
keyDict = keySplits[1] + '/' + keySplits[1] + '/' + keySplits[2]
f[keyDict] = reader.get_tensor(key)
以上這篇tensorflow 模型權重導出實例就是小編分享給大家的全部內容了,希望能給大家一個參考,也希望大家多多支持腳本之家。
相關文章
keras Lambda自定義層實現數據的切片方式,Lambda傳參數
這篇文章主要介紹了keras Lambda自定義層實現數據的切片方式,Lambda傳參數,具有很好的參考價值,希望對大家有所幫助。一起跟隨小編過來看看吧2020-06-06

