TensorFlow Estimator图冻结失败：Iris示例转Core ML遇错误

嘿，我之前折腾过把Iris Estimator转成CoreML模型，也踩过冻结图导出失败的坑——主要是Estimator的模型结构和传统Session模型不一样，直接用graph.pbtxt和model.ckpt硬转肯定出问题。下面是我亲测有效的完整流程，一步步来应该能解决你的问题：

第一步：先导出Estimator的SavedModel格式（关键！） #

别再盯着graph.pbtxt和model.ckpt了，Estimator官方推荐用export_saved_model导出完整的SavedModel，这能避免很多内部graph结构的坑。给你贴个适配Iris模型的示例代码：

import tensorflow as tf
from tensorflow.contrib.learn.python.learn.datasets import load_iris

# 定义你的Iris Estimator模型函数（如果已经训练好可以跳过训练部分）
def iris_model_fn(features, labels, mode):
    # 模型结构：4维输入→全连接层→3类输出
    logits = tf.layers.dense(features["x"], 3, activation=None)
    predictions = tf.argmax(logits, axis=1, name="iris_prediction")

    # 预测模式的返回（这部分是导出的核心）
    if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode, predictions=predictions)

    # 训练&评估部分（如果还没训练就执行这部分）
    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
    train_op = tf.train.AdamOptimizer(learning_rate=0.01).minimize(
        loss, global_step=tf.train.get_global_step())
    return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)

# 加载Iris数据集
iris_data = load_iris()
train_input_fn = tf.estimator.inputs.numpy_input_fn(
    x={"x": iris_data.data}, y=iris_data.target, batch_size=8, shuffle=True, num_epochs=None)

# 创建Estimator并训练（如果已训练好可跳过）
estimator = tf.estimator.Estimator(model_fn=iris_model_fn, model_dir="./iris_trained_model")
estimator.train(input_fn=train_input_fn, steps=1000)

# 导出SavedModel：定义输入特征格式，匹配Iris的4维输入
feature_spec = {"x": tf.FixedLenFeature(shape=(4,), dtype=tf.float32)}
serving_receiver_fn = tf.estimator.export.build_parsing_serving_input_receiver_fn(feature_spec)
export_path = estimator.export_saved_model("./iris_saved_model", serving_receiver_fn)

注意：serving_receiver_fn必须和你的输入特征维度完全匹配（Iris是4维特征），这一步是告诉TensorFlow模型在预测时接受什么样的输入格式。

第二步：用SavedModel生成冻结图 #

有了SavedModel，再用官方的freeze_graph.py工具生成冻结图，这比直接用ckpt文件靠谱太多。命令如下：

python tensorflow/python/tools/freeze_graph.py \
  --input_saved_model_dir=./iris_saved_model/[你的SavedModel版本文件夹] \
  --output_graph=./frozen_iris.pb \
  --output_node_names="iris_prediction"

找不到输出节点名？用TensorBoard看SavedModel的graph结构就行：

tensorboard --logdir=./iris_saved_model

打开Graph标签，找到你在模型里定义的预测输出节点（比如上面代码里的iris_prediction），把它的名字填到--output_node_names里。

第三步：冻结图转CoreML模型 #

用tfcoreml工具完成最后一步转换，命令示例：

tfcoreml convert \
  --tf-model-path=./frozen_iris.pb \
  --mlmodel-output=./IrisClassifier.mlmodel \
  --input-feature-names="x:0" \
  --output-feature-names="iris_prediction:0"

或者用Python代码转换（更灵活）：

import tfcoreml

tfcoreml.convert(
    tf_model_path="./frozen_iris.pb",
    mlmodel_path="./IrisClassifier.mlmodel",
    input_name_shape_dict={"x:0": (1, 4)},  # 1是batch size，4是Iris特征数
    output_feature_names=["iris_prediction:0"]
)

常见坑排查 #

• 冻结图时提示节点不存在：检查--output_node_names的拼写，或者确认你的Estimator在PREDICT模式下正确返回了预测节点（别漏写mode == tf.estimator.ModeKeys.PREDICT的分支）。
• CoreML转换时输入形状不匹配：确保input_name_shape_dict的形状和模型输入一致，Iris是4维特征，所以填(1,4)（1代表单样本预测）。
• SavedModel导出失败：检查你的模型训练时的model_dir是否正确，有没有生成合法的ckpt文件。

如果还是有具体的错误信息，可以贴出来，我再帮你针对性分析~

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