加载tf.contrib.factorization.KMeansClustering模型时触发KeyError: 'NearestNeighbors'

这个错误的核心原因是：NearestNeighbors是tf.contrib.factorization模块下的自定义操作符，而TensorFlow 1.x不会自动注册contrib模块里的所有op——如果加载模型的脚本没有显式导入该模块，TensorFlow就找不到这个操作符。

下面是两种可行的解决方法：

方法1：显式导入contrib.factorization模块 #

在加载模型的代码开头，添加导入语句，强制TensorFlow注册NearestNeighbors操作符：

import tensorflow as tf
# 显式导入factorization模块，确保相关op被注册
from tensorflow.contrib import factorization

with tf.Session() as sess:
    tf.saved_model.loader.load(
        sess, 
        [tf.saved_model.tag_constants.SERVING], 
        "/path/to/your/saved_model"
    )

这是最直接的解决方案，因为导入模块时TensorFlow会自动注册其中定义的所有自定义操作符。

方法2：确保导出模型的计算图完整性 #

有时候如果KMeans模型未完成至少一次训练就导出，计算图中可能缺少NearestNeighbors相关的节点。建议你在导出前先运行至少一轮训练（哪怕是小批量数据），确保所有必要的op都被添加到计算图中。

另外，你也可以调整serving_input_receiver_fn，让输入处理逻辑和训练时更一致，避免不必要的图结构差异：

def serving_input_receiver_fn():
    # 直接使用和训练时一致的输入占位符，避免parse_example带来的额外图节点
    x = tf.placeholder(tf.float32, shape=[None, 2], name="input_x")
    receiver_tensors = {"x": x}
    features = {"x": x}
    return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)

修改后重新导出模型，再尝试加载。

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