Keras模型转换为TensorFlow二进制及Serving部署指南

花开有时
花开有时
发布时间:2026-10-13| 7人浏览过
原创

本文详解如何将Keras训练的模型转换为TensorFlow二进制格式(.pb),并解决直接转换导致的TensorFlow Serving加载失败问题。通过冻结会话生成常量图,再进一步转换为SavedModel格式,确保模型可在服务端正常推理部署。

Keras模型转换为TensorFlow二进制格式

近期需要将基于Keras训练的模型移植至移动端运行,因此需将其转换为TensorFlow的二进制模型格式。

经过多次尝试,找到了一种可行的转换方案,以下直接展示具体实现代码。


# coding=utf-8
import sys

from keras.models import load_model
import tensorflow as tf
import os
import os.path as osp
from keras import backend as K

def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True):
 """
 Freezes the state of a session into a prunned 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
 prunned so subgraphs that are not neccesary 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()]
  input_graph_def = graph.as_graph_def()
  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

input_fld = sys.path[0]
weight_file = 'your_model.h5'
output_graph_name = 'tensor_model.pb'

output_fld = input_fld + '/tensorflow_model/'
if not os.path.isdir(output_fld):
 os.mkdir(output_fld)
weight_file_path = osp.join(input_fld, weight_file)

K.set_learning_phase(0)
net_model = load_model(weight_file_path)

print('input is :', net_model.input.name)
print ('output is:', net_model.output.name)

sess = K.get_session()

frozen_graph = freeze_session(K.get_session(), output_names=[net_model.output.op.name])

from tensorflow.python.framework import graph_io

graph_io.write_graph(frozen_graph, output_fld, output_graph_name, as_text=False)

print('saved the constant graph (ready for inference) at: ', osp.join(output_fld, output_graph_name))

上述代码执行后,模型文件将保存至当前目录下的tensor_model文件夹中。

接下来对转换后的模型进行验证测试。


import tensorflow as tf
import numpy as np
import PIL.Image as Image
import cv2

def recognize(jpg_path, pb_file_path):
 with tf.Graph().as_default():
  output_graph_def = tf.GraphDef()

  with open(pb_file_path, "rb") as f:
   output_graph_def.ParseFromString(f.read())
   tensors = tf.import_graph_def(output_graph_def, name="")
   print tensors

  with tf.Session() as sess:
   init = tf.global_variables_initializer()
   sess.run(init)

   op = sess.graph.get_operations()
   
   for m in op:
    print(m.values())

   input_x = sess.graph.get_tensor_by_name("convolution2d_1_input:0") #具体名称看上一段代码的input.name
   print input_x

   out_softmax = sess.graph.get_tensor_by_name("activation_4/Softmax:0") #具体名称看上一段代码的output.name

   print out_softmax

   img = cv2.imread(jpg_path, 0)
   img_out_softmax = sess.run(out_softmax,
          feed_dict={input_x: 1.0 - np.array(img).reshape((-1,28, 28, 1)) / 255.0})

   print "img_out_softmax:", img_out_softmax
   prediction_labels = np.argmax(img_out_softmax, axis=1)
   print "label:", prediction_labels

pb_path = 'tensorflow_model/constant_graph_weights.pb'
img = 'test/6/8_48.jpg'
recognize(img, pb_path)

解决TensorFlow Serving部署报错问题

补充说明:如何将Keras训练好的模型转换为TensorFlow的.pb文件,并在TensorFlow Serving环境中进行调用。

首先,Keras训练完成的模型通过内置的model.save()方法保存时,默认生成的是.h5格式的文件。

在Keras中载入模型的标准方式为:my_model = keras.models.load_model(filepath)。

若需将该模型转换为.pb格式的TensorFlow模型,可参考以下代码实现。


# -*- coding: utf-8 -*-
from keras.layers.core import Activation, Dense, Flatten
from keras.layers.embeddings import Embedding
from keras.layers.recurrent import LSTM
from keras.layers import Dropout
from keras.layers.wrappers import Bidirectional
from keras.models import Sequential,load_model
from keras.preprocessing import sequence
from sklearn.model_selection import train_test_split
import collections
from collections import defaultdict
import jieba
import numpy as np
import sys
reload(sys)
sys.setdefaultencoding('utf-8')
import tensorflow as tf
import os
import os.path as osp
from keras import backend as K
def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True):
 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()]
  input_graph_def = graph.as_graph_def()
  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
input_fld = '/data/codebase/Keyword-fenci/brand_recogniton_biLSTM/'
weight_file = 'biLSTM_brand_recognize.model'
output_graph_name = 'tensor_model_v3.pb'

output_fld = input_fld + '/tensorflow_model/'
if not os.path.isdir(output_fld):
 os.mkdir(output_fld)
weight_file_path = osp.join(input_fld, weight_file)

K.set_learning_phase(0)
net_model = load_model(weight_file_path)

print('input is :', net_model.input.name)
print ('output is:', net_model.output.name)

sess = K.get_session()

frozen_graph = freeze_session(K.get_session(), output_names=[net_model.output.op.name])
from tensorflow.python.framework import graph_io

graph_io.write_graph(frozen_graph, output_fld, output_graph_name, as_text=True)

print('saved the constant graph (ready for inference) at: ', osp.join(output_fld, output_graph_name))

执行上述代码后,模型即被保存为.pb格式的文件。

此时产生的问题是,这样生成的.pb格式文件属于frozen model(冻结模型)。

若直接通过TensorFlow Serving加载该模型,将会遇到报错信息。

错误日志显示:E tensorflow_serving/core/aspired_versions_manager.cc:358] Servable {name: mnist version: 1} cannot be loaded: Not found: Could not find meta graph def matching supplied tags: { serve }. To inspect available tag-sets in the SavedModel, please use the SavedModel CLI: `saved_model_cli`。

这是因为TensorFlow Serving期望读取的是SavedModel格式,而非冻结图。

因此,需要将frozen model进一步转换为SavedModel格式,具体解决方案如下。


from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants

export_dir = '/data/codebase/Keyword-fenci/brand_recogniton_biLSTM/saved_model'
graph_pb = '/data/codebase/Keyword-fenci/brand_recogniton_biLSTM/tensorflow_model/tensor_model.pb'

builder = tf.saved_model.builder.SavedModelBuilder(export_dir)

with tf.gfile.GFile(graph_pb, "rb") as f:
 graph_def = tf.GraphDef()
 graph_def.ParseFromString(f.read())

sigs = {}

with tf.Session(graph=tf.Graph()) as sess:
 # name="" is important to ensure we don't get spurious prefixing
 tf.import_graph_def(graph_def, name="")
 g = tf.get_default_graph()
 inp = g.get_tensor_by_name(net_model.input.name)
 out = g.get_tensor_by_name(net_model.output.name)

 sigs[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = \
  tf.saved_model.signature_def_utils.predict_signature_def(
   {"in": inp}, {"out": out})

 builder.add_meta_graph_and_variables(sess,
           [tag_constants.SERVING],
           signature_def_map=sigs)
builder.save()

转换完成后,保存的SavedModel文件夹中将包含以下两个部分。

包括saved_model.pb文件和variables文件夹。

其中variables文件夹可以为空,不影响模型加载。

将转换后的.pb模型导入Serving并读取,即可成功运行。

以上关于Keras模型保存为TensorFlow二进制模型的方法即为全部内容,希望能为读者提供参考,感谢支持。

免责声明:本文内容仅供参考,不代表本站观点。转载或引用请注明出处。

Python 相关产品和应用

全部下载
Selenium

Selenium

4.51.0 · Windows

用于浏览器自动化和 Web 测试的 Python 工具。

scikit-learn

scikit-learn

1.9.1 · Windows

常用 Python 机器学习算法与模型评估工具库。

Black

Black

26.10.1 · Windows

强调一致性的 Python 自动代码格式化工具。

Ruff

Ruff

0.17.1 · macOS

由 Rust 编写的高速 Python 代码检查和格式化工具。

pytest

pytest

9.1.1 · Windows

简洁而强大的 Python 自动化测试框架。

FastAPI

FastAPI

0.146.0 · Windows

高性能 Python API 框架,支持类型提示和自动接口文档。

Flask

Flask

3.1.3 · Windows

轻量灵活的 Python Web 应用框架。

Django

Django

6.1.2 · Windows

成熟的 Python Web 框架,适合构建内容站点和业务系统。

Spyder

Spyder

6.1.7 · Windows

面向科学计算与数据分析的 Python 集成开发环境。

SciPy

SciPy

1.18.1 · Windows

面向科学、工程和技术计算的 Python 算法库。