TensorFlow 模型转换 .pb 转 .lite 实例教程

远山来客

发布时间：2026-10-13| 6人浏览过

原创

本文详细介绍如何将现有的 TensorFlow .pb 分类模型转换为适用于移动端的 .lite 格式。通过编译 TensorFlow 源码获取 TOCO 工具，并使用命令行参数指定输入输出节点、形状及数据类型完成转换。文章还提供了查找未知节点的方法，并对比了转换前后模型在 Android 设备上的大小与推理速度，展示 Lite 格式的加速效果。

TensorFlow Lite 模型转换背景

2017年底，TensorFlow 推出了 Lite 版本，旨在实现移动端的快速运行。其中，一个关键问题是如何将现有的分类模型（.pb）转换为 .lite 模型。

具体的转换步骤如下所示。

第一步，进入 TensorFlow 源码文件夹。这样做是为了让 bazel 能够无需额外配置即可找到相应的路径。

PB 转 Lite 转换步骤

第二步，生成转换所需的可执行文件。


bazel run --config=opt \
 //tensorflow/contrib/lite/toco:toco -- \
 --input_file=/tmp/mobilenet_v1_1.0_224_frozen.pb \  # 待转换模型路径
 --output_file=/tmp/tflite_model2.tflite \      # 目标模型路径
 --input_format=TENSORFLOW_GRAPHDEF \
 --output_format=TFLITE \
 --input_shape=1,224,224,3 \              # 输入图像宽高
 --input_array=input \                 # 输入节点名称<参考1>
 --output_array=MobilenetV1/Predictions/Reshape_1 \  # 输出节点名称<参考1>
 --inference_type=FLOAT \               # 图像数据类型
 --input_data_type=FLOAT

查找输入输出节点方法

参考1：如果不知道输入和输出节点，可以进入 Python 环境，导入 tf 模块，并通过以下命令（tf.GraphDef()）进行查找：


>>> import tensorflow as tf
>>> gf = tf.GraphDef()
>>> gf.ParseFromString(open('/your/path/to/graphname.pb','rb').read())
>>> for n in gf.node:
>>>   print ( n.name +' ===> '+n.op )