在制作TensorFlow Lite模型时如何设置input_array与output_array名称

结合你用的Ubuntu 14.04、TensorFlow 1.4.0（二进制安装）、CUDA 8.0这些环境，我来一步步帮你完成模型转换：

第一步：导出冻结图（Frozen Graph） #

TF1.4.0的TFLite Converter对冻结图的支持最成熟，所以先把训练好的模型转成冻结图格式：

1. 先保存模型的图结构：
   在你的训练代码末尾加上这段，把图结构写成文本格式的pbtxt文件：

   import tensorflow as tf
   from tensorflow.python.tools.freeze_graph import freeze_graph
   
   # 获取训练好的模型图
   graph = tf.get_default_graph()
   # 把图结构写入文件
   tf.train.write_graph(graph.as_graph_def(), "./", "p_net_graph.pbtxt", as_text=True)
   

2. 冻结图（把权重和图结构合并成一个pb文件）：
   你需要先确定模型的输入节点名（就是你P_Net函数里的inputs对应的节点名，一般就是"inputs"）和输出节点名（比如分类输出、bbox预测输出、landmark预测输出，你得根据自己的模型定义找对应的节点名称，比如可能是p_net/softmax/Softmax、p_net/bbox_pred/BiasAdd、p_net/landmark_pred/BiasAdd）。然后运行下面的代码：

   # 定义输出节点列表，替换成你实际的节点名
   output_nodes = ["p_net/softmax/Softmax", "p_net/bbox_pred/BiasAdd", "p_net/landmark_pred/BiasAdd"]
   
   freeze_graph(
       input_graph="./p_net_graph.pbtxt",
       input_saver="",
       input_binary=False,
       input_checkpoint="./your_trained_model.ckpt", # 替换成你训练时保存的ckpt文件路径
       output_node_names=",".join(output_nodes),
       restore_op_name="save/restore_all",
       filename_tensor_name="save/Const:0",
       output_graph="./frozen_p_net.pb", # 输出的冻结图文件
       clear_devices=True,
       initializer_nodes=""
   )
   

第二步：用TFLite Converter转换冻结图 #

TF1.4.0里的Converter在tf.contrib.lite模块下，直接用它来转冻结图：

import tensorflow as tf

# 初始化转换器
converter = tf.contrib.lite.TFLiteConverter.from_frozen_graph(
    graph_file="./frozen_p_net.pb",
    input_arrays=["inputs"], # 你的模型输入节点名
    input_shapes={"inputs": [1, None, None, 3]}, # 输入形状，这里设成[1, None, None, 3]支持任意尺寸的输入图片，根据你训练时的输入调整
    output_arrays=output_nodes # 之前定义的输出节点列表
)

# 因为你的模型用了PReLU激活函数，TF1.4.0的TFLite默认可能不支持，所以开启自定义操作允许选项
converter.allow_custom_ops = True

# 执行转换
tflite_model = converter.convert()

# 保存转换好的TFLite模型
with open("./p_net.tflite", "wb") as f:
    f.write(tflite_model)

第三步：可能踩的坑及解决办法 #

• PReLU不被识别：如果转换时报错说PReLU是未支持的操作，最简单的办法是把模型里的PReLU替换成普通ReLU重新训练一轮；如果不想重训，也可以自定义PReLU的TFLite操作，但这个步骤比较繁琐，适合有一定经验的开发者。
• 输入尺寸固定问题：如果你的模型是用于人脸检测这类需要可变输入尺寸的场景，一定要把input_shapes里的宽高维度设为None，不然转换后的模型只能处理固定尺寸的图片。
• CUDA环境干扰：转换过程不需要用到CUDA，你可以在纯CPU环境下操作，避免CUDA版本带来的依赖问题。

第四步：验证转换后的模型 #

转换完成后，用TF Lite的Python接口快速验证模型是否能正常运行：

import tensorflow as tf
import numpy as np

# 加载TFLite模型
interpreter = tf.contrib.lite.Interpreter(model_path="./p_net.tflite")
interpreter.allocate_tensors()

# 获取输入输出张量的详情
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 生成一张测试用的随机图片（和你训练时的输入格式一致，比如float32类型）
test_input = np.random.rand(1, 128, 128, 3).astype(np.float32)
# 设置输入张量
interpreter.set_tensor(input_details[0]['index'], test_input)

# 运行推理
interpreter.invoke()

# 获取输出结果
output_class = interpreter.get_tensor(output_details[0]['index'])
output_bbox = interpreter.get_tensor(output_details[1]['index'])
output_landmark = interpreter.get_tensor(output_details[2]['index'])

# 打印输出形状，确认和原模型一致
print(f"分类输出形状：{output_class.shape}")
print(f"BBox预测输出形状：{output_bbox.shape}")
print(f"Landmark预测输出形状：{output_landmark.shape}")

内容的提问来源于stack exchange，提问作者flankechen