1. 程式人生 > 程式設計 >tensorflow 2.0模式下訓練的模型轉成 tf1.x 版本的pb模型例項

tensorflow 2.0模式下訓練的模型轉成 tf1.x 版本的pb模型例項

升級到tf 2.0後,訓練的模型想轉成1.x版本的.pb模型,但之前提供的通過ckpt轉pb模型的方法都不可用(因為儲存的ckpt不再有.meta)檔案,嘗試了好久,終於找到了一個方法可以迂迴轉到1.x版本的pb模型.

Note: 本方法首先有些要求需要滿足:

可以拿的到模型的網路結構定義原始碼

網路結構裡面的所有操作都是通過tf.keras完成的,不能出現類似tf.nn 的tensorflow自己的操作符

tf2.0下儲存的模型是.h5格式的,並且僅儲存了weights,即通過model.save_weights儲存的模型.

在tf1.x的環境下,將tf2.0儲存的weights轉為pb模型:

如果在tf2.0下儲存的模型符合上述的三個定義,那麼這個.h5檔案在1.x環境下其實是可以直接用的,因為都是通過tf.keras高階封裝了,2.0版本和1.x版本不存在特別大的區別,我自己的模型是可以直接用的.

import tensorflow as tf
import os
from nets.efficientNet import *
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
# 這個程式碼網上說需要加上,如果模型裡有dropout,bn層的話,我測試過加不加結果都一樣,保險起見還是加上吧
tf.keras.backend.set_learning_phase(0)

# 首先是定義你的模型,這個需要和tf2.0下一毛一樣
inputs = tf.keras.Input(shape=(224,224,3),name='modelInput')
outputs = yourModel(inputs,training=False)
model = tf.keras.Model(inputs=inputs,outputs=outputs)
model.load_weights('save_weights.h5')
def freeze_session(session,keep_var_names=None,output_names=None,clear_devices=True):
  """
  Freezes the state of a session into a pruned computation graph.

  Creates a new computation graph where variable nodes are replaced by
  constants taking their current value in the session. The new graph will be
  pruned so subgraphs that are not necessary to compute the requested
  outputs are removed.
  @param session The TensorFlow session to be frozen.
  @param keep_var_names A list of variable names that should not be frozen,or None to freeze all the variables in the graph.
  @param output_names Names of the relevant graph outputs.
  @param clear_devices Remove the device directives from the graph for better portability.
  @return The frozen graph definition.
  """
  from tensorflow.python.framework.graph_util import convert_variables_to_constants
  graph = session.graph
  with graph.as_default():
    freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference(keep_var_names or []))
    output_names = output_names or []
    output_names += [v.op.name for v in tf.global_variables()]
    # Graph -> GraphDef ProtoBuf
    input_graph_def = graph.as_graph_def(add_shapes=True)
    if clear_devices:
      for node in input_graph_def.node:
        node.device = ""
    frozen_graph = convert_variables_to_constants(session,input_graph_def,output_names,freeze_var_names)
    return frozen_graph

frozen_graph = freeze_session(tf.keras.backend.get_session(),output_names=[out.op.name for out in model.outputs])
tf.train.write_graph(frozen_graph,"model","tf_model.pb",as_text=False)

執行成功後,會在當前目錄下生成一個model資料夾,裡面有生成的tf_model.pb檔案,至此,我們就完成了將tf2.0下訓練的模型轉到tf1.x下的pb模型,這樣,就可以用這個pb模型做其它推理或者轉tvm ncnn等模型轉換工作.

這個轉換的重點就是通過keras這個中間商來完成,所以我們定義的模型就必須要滿足這個中間商定義的條件

補充知識:tensorflow2.0降級及如何從別的版本升到2.0

程式碼實踐《tensorflow實戰GOOGLE深度學習框架》時,由於本機安裝的tensorflow為2.0版本與配套書籍程式碼1.4的API不相容,只得將tensorflow降級為1.4.0版本使用,降級方法如下

1 pip uninstall tensorflow

tensorflow 2.0模式下訓練的模型轉成 tf1.x 版本的pb模型例項

2 pip install tensorflow==1.14.0 -i https://pypi.tuna.tsinghua.edu.cn/simple

tensorflow 2.0模式下訓練的模型轉成 tf1.x 版本的pb模型例項

驗證

import tensorflow as tf
print(tf.version)

tensorflow 2.0模式下訓練的模型轉成 tf1.x 版本的pb模型例項

二 從別的版本升級到2.0

自動解除安裝與其相關包

pip uninstall tensorflow

安裝某版本

pip install --no-cache-dir tensorflow==x.xx (此處填寫2.0)

tensorflow 2.0模式下訓練的模型轉成 tf1.x 版本的pb模型例項

驗證

tensorflow 2.0模式下訓練的模型轉成 tf1.x 版本的pb模型例項

以上這篇tensorflow 2.0模式下訓練的模型轉成 tf1.x 版本的pb模型例項就是小編分享給大家的全部內容了,希望能給大家一個參考,也希望大家多多支援我們。