Python冻结Keras模型并在Go中加载及性能优化相关问题

我来帮你逐个梳理这些TensorFlow跨语言部署的常见问题，都是实际踩过的坑：

1. 正确冻结模型并在Go中加载的方法 #

首先要明确：你用tf.LoadSavedModel加载的是SavedModel格式的模型，但你当前生成的frozen_my_model.pb是GraphDef格式的冻结模型，这俩不是同一种格式，所以才会出现找不到serve标签的错误。这里有两种靠谱的解决路径：

路径一：直接导出SavedModel（最简单） #

跳过手动冻结的步骤，训练完Keras模型后直接导出成SavedModel格式，自带serve标签，Go端可以直接加载：

# Python训练完模型后执行
m1.save("my_saved_model", save_format="tf")

这会生成一个my_saved_model文件夹，里面包含SavedModel的标准结构。然后Go端加载这个文件夹路径即可：

model, err := tf.LoadSavedModel("my_saved_model", []string{"serve"}, nil)
if err != nil {
    // 处理错误
}

路径二：加载冻结的GraphDef模型（如果坚持用冻结.pb） #

如果一定要用冻结后的.pb文件，Go端不能用LoadSavedModel，要直接导入GraphDef并创建Session：

import (
    "io/ioutil"
    "github.com/tensorflow/tensorflow/tensorflow/go"
)

func loadFrozenModel() (*tensorflow.Session, *tensorflow.Graph, error) {
    // 读取冻结模型文件
    modelBytes, err := ioutil.ReadFile("frozen_my_model.pb")
    if err != nil {
        return nil, nil, err
    }

    // 导入GraphDef
    graph := tensorflow.NewGraph()
    if err := graph.Import(modelBytes, ""); err != nil {
        return nil, nil, err
    }

    // 创建推理Session
    sess, err := tensorflow.NewSession(graph, nil)
    if err != nil {
        return nil, nil, err
    }

    return sess, graph, nil
}

关键注意点：冻结模型时一定要确认output_node_names的正确性，你可以在Python里打印模型输出节点的名称：

print(m1.output.name)  # 比如会输出"outputNode/Softmax:0"，冻结时要去掉末尾的":0"，用"outputNode/Softmax"

2. 冻结模型是否能提升推理速度？ #

是的，冻结模型确实能有效提升推理速度，原因有两个：

• 冻结过程会把所有可训练变量转换成常量，避免了推理时的变量初始化、查找等额外开销；
• 冻结后可以配合优化工具移除所有训练相关的节点（比如梯度计算、Dropout的训练分支、学习率更新等），让模型更轻量化，减少内存占用和推理延迟。

对于Go这种追求性能的部署场景，冻结+优化后的模型会比未冻结的模型表现好很多。

3. 如何加入optimize_for_inference优化？ #

optimize_for_inference是TensorFlow提供的推理优化工具，能移除训练节点、合并冗余操作，进一步压缩模型并提升速度。你可以在冻结模型之后执行以下步骤：

from tensorflow.python.tools import optimize_for_inference_lib
import tensorflow as tf

# 读取冻结后的模型
with tf.gfile.GFile("frozen_my_model.pb", "rb") as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())

# 替换成你模型的实际输入/输出节点名（注意要和冻结时的一致）
input_node_names = ["inputNode_input"]  # 可以用print(m1.input.name)查看，去掉末尾的":0"
output_node_names = ["outputNode/Softmax"]

# 执行优化
optimized_graph_def = optimize_for_inference_lib.optimize_for_inference(
    graph_def,
    input_node_names,
    output_node_names,
    tf.float32.as_datatype_enum  # 替换成你模型输入的数据类型，比如tf.int32.as_datatype_enum
)

# 保存优化后的模型
with tf.gfile.GFile("optimized_frozen_my_model.pb", "wb") as f:
    f.write(optimized_graph_def.SerializeToString())

优化后的模型可以直接用路径二中的Go代码加载，推理速度会比单纯冻结的模型更快。

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