Internal change

PiperOrigin-RevId: 460979064

Internal change
PiperOrigin-RevId: 460979064
2c4ea3d8 · Yang Feng · A. Unique TensorFlower · c9e4db9c · 2c4ea3d8
Commit 2c4ea3d8 authored Jul 14, 2022 by Yang Feng Committed by A. Unique TensorFlower Jul 14, 2022
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

official/vision/serving/image_classification.py official/vision/serving/image_classification.py +4 -1

No files found.
--- a/official/vision/serving/image_classification.py
+++ b/official/vision/serving/image_classification.py
@@ -78,6 +78,9 @@ class ClassificationModule(export_base.ExportModule):
                parallel_iterations=32))

    logits = self.inference_step(images)
+    if self.params.task.train_data.is_multilabel:
+      probs = tf.math.sigmoid(logits)
+    else:
      probs = tf.nn.softmax(logits)

    return {'logits': logits, 'probs': probs}