CODEWITHSUNDEEP
tensorflowtensorflow-servingtensorflow-estimator
Byest
Byest

Reputation: 337

In Tensorflow-Serving, is it possible to get only the top-k prediction results?

When using the code in https://www.tensorflow.org/serving, but with a DNNClassifier Estimator model, the curl/query request returns all the possible label classes and their associated scores.

Using a model with 100,000+ possible output/label classes, the response becomes too large. Is there any way to limit the number of outputs to the top-k results? (Similar to how it can be done in keras).

The only possibility I could think of is feeding some parameter into the predict API through the signatures, but I haven't found any parameters that would give this functionality. I've read through a ton of documentation + code and googled a ton, but to no avail.

Any help would be greatly appreciated. Thanks in advance for any responses. <3

Upvotes: 2

Views: 519

Answers (2)

zzachimonde
zzachimonde

Reputation: 454

AFAIC, there are 2 ways to support your need.

  1. You could add some lines in tensorflow-serving source code referring to this

  2. You could do something like this while training/retraining your model.

Hope this will help.

Upvotes: 2

Byest
Byest

Reputation: 337

Putting this up here in case it helps anyone. It's possible to override the classification_output() function in head.py (which is used by dnn.py) in order to filter the top-k results. You can insert this snippet into your main.py / train.py file, and whenever you save an DNNClassifier model, that model will always output at most num_top_k_results when doing inference/serving. The vast majority of the method is copied from the original classification_output() function. (Note this may or may not work with 1.13 / 2.0 as it hasn't been tested on those.)

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

Upvotes: 0

Related Questions