Keras模型转换为TensorFlow .pb及TFLite部署指南

Keras模型转换为TensorFlow .pb格式

在模型部署过程中，将Keras的 .h5 格式模型转换为TensorFlow的 .pb 格式是常见需求。这种格式便于在前端环境中进行推理部署，同时也为生成移动端专用的 .tflite 文件奠定了基础。以下操作基于 tensorflow 1.12.0 和 python 3.6.5 环境。

1. 保存Keras模型

首先，确保你的Keras模型已正确保存为 .h5 格式。可以使用以下代码保存模型：

model.save('tf.h5')

2. 转换代码实现

接下来，通过Python脚本将 .h5 模型转换为 .pb 文件。关键步骤包括指定输入和输出变量名，并冻结计算图。

from keras.models import Model
from keras.layers import Dense, Dropout
from keras.applications.mobilenet import MobileNet
from keras.applications.mobilenet import preprocess_input
from keras.preprocessing.image import load_img, img_to_array
import tensorflow as tf
from keras import backend as K
import os
 
# 构建模型
base_model = MobileNet((None, None, 3), alpha=1, include_top=False, pooling='avg', weights=None)
x = Dropout(0.75)(base_model.output)
x = Dense(10, activation='softmax')(x)
 
model = Model(base_model.input, x)
model.load_weights('mobilenet_weights.h5')
 
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
 
output_graph_name = 'NIMA.pb'
output_fld = ''
# K.set_learning_phase(0)
 
print('input is :', model.input.name)
print ('output is:', model.output.name)
 
sess = K.get_session()
frozen_graph = freeze_session(K.get_session(), output_names=[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: ', os.path.join(output_fld, output_graph_name))

3. 关键参数说明

• input_arrays：指定模型起始输入的变量名。
• output_arrays：指定模型结束输出的变量名。
• input_shape：与 input_arrays 对应的输入形状。

TFLite模型转换

对于移动端部署，可以将 .pb 文件进一步转换为 .tflite 格式。使用以下命令进行转换：

tflite_convert --output_file=tf.tflite --keras_model_file=tf.h5

如果在转换过程中遇到 TypeError: __init__() missing 2 required positional arguments: 'filters' and 'kernel_size' 错误，请检查模型定义是否正确。

常见报错与排查

在转换过程中，可能会遇到 convolution2d_1_input 找不到的报错。此时，可以通过以下步骤进行排查：

1. 在堆栈中找到 convert_saved_model.py 文件。
2. 定位到 get_tensors_from_tensor_names() 方法。
3. 添加 print(list(tensor_name_to_tensor)) 以打印所有tensor的名字，从而判断实际的名称。

使用keras_to_tensorflow工具

此外，还可以使用 keras_to_tensorflow 工具进行转换。具体步骤如下：

git clone git@github.com:amir-abdi/keras_to_tensorflow.git
cd keras_to_tensorflow
python keras_to_tensorflow.py --input_model=path/to/tf.h5 --output_model=path/to/tf.pb
tflite_convert \
    --output_file=tf.tflite \
    --graph_def_file=tf.pb \
    --input_arrays=convolution2d_1_input \
    --output_arrays=dense_3/BiasAdd \
    --input_shape=1,3,448,448

通过上述步骤，您可以成功将Keras模型转换为TensorFlow的 .pb 格式，并进一步转换为 .tflite 格式，以满足不同平台的部署需求。

以上就是Keras模型转成tensorflow的.pb操作的详细内容，更多关于Keras模型转换的资料请关注本站其它相关文章！