怎么将tensorflow2.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下保存的模型是.h6格式的,并且仅保存了weights, 即通过model.save_weights保存的模型.

在tf1.x的环境下, 将tf2.0保存的weights转为pb模型:

如果在tf2.0下保存的模型符合上述的三个定义, 那么这个.h6文件在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.h6')
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)

另外有需要云服务器可以了解下创新互联scvps.cn，海内外云服务器15元起步，三天无理由+7*72小时售后在线，公司持有idc许可证，提供“云服务器、裸金属服务器、高防服务器、香港服务器、美国服务器、虚拟主机、免备案服务器”等云主机租用服务以及企业上云的综合解决方案，具有“安全稳定、简单易用、服务可用性高、性价比高”等特点与优势，专为企业上云打造定制，能够满足用户丰富、多元化的应用场景需求。