在Tensorflow-Serving中，是否可以仅获得前k个预测结果？

Question

当使用https://www.tensorflow.org/serving中的代码但使用DNNClassifier Estimator模型时，curl / query请求将返回所有可能的标签类及其相关分数。

使用具有100,000多个可能的输出/标签类的模型，响应变得太大。有没有办法限制前k个结果的输出数量？ （类似于如何在keras中完成）。

我能想到的唯一可能性是通过签名将一些参数输入到预测API中，但我没有找到任何可以提供此功能的参数。我已经阅读了大量文档+代码并用谷歌搜索了一下，但无济于事。

任何帮助将不胜感激。提前感谢您的回复。 <3

Answer 1

AFAIC，有两种方式可以满足您的需求。

1. 您可以在tensorflow服务源代码中添加一些引用this的行
2. 在训练/重新训练你的模型时，你可以做像this这样的事情。

希望这会有所帮助。

Answer 2

把它放在这里以防万一它可以帮助任何人。可以覆盖head.py（由dnn.py使用）中的classification_output（）函数，以过滤top-k结果。您可以将此代码段插入到main.py / train.py文件中，并且每当您保存DNNC分类器模型时，该模型在进行推理/提供时始终最多输出num_top_k_results。绝大多数方法都是从原始的classification_output（）函数复制而来的。 （请注意，这可能适用于1.13 / 2.0，也可能不适用，因为它尚未经过测试。）

from tensorflow.python.estimator.canned import head as head_lib

num_top_k_results = 5

def override_classification_output(scores, n_classes, label_vocabulary=None):
  batch_size = array_ops.shape(scores)[0]
  if label_vocabulary:
    export_class_list = label_vocabulary
  else:
    export_class_list = string_ops.as_string(math_ops.range(n_classes))
  # Get the top_k results
  top_k_scores, top_k_indices = tf.nn.top_k(scores, num_top_k_results)
  # Using the top_k_indices, get the associated class names (from the vocabulary)
  top_k_classes = tf.gather(tf.convert_to_tensor(value=export_class_list), tf.squeeze(top_k_indices))
  export_output_classes = array_ops.tile(
      input=array_ops.expand_dims(input=top_k_classes, axis=0),
      multiples=[batch_size, 1])
  return export_output.ClassificationOutput(
      scores=top_k_scores,
      # `ClassificationOutput` requires string classes.
      classes=export_output_classes)

# Override the original method with our custom one.
head_lib._classification_output = override_classification_output