tf.math.top_k 不接受整数 {"error": "k 必须是标量，形状为 [1]\n\t [[{{node TopKV2}}]]"}答案

【问题标题】：tf.math.top_k not accepting integers {"error": "k must be scalar, got shape [1]\n\t [[{{node TopKV2}}]]"}tf.math.top_k 不接受整数 {"error": "k 必须是标量，形状为 [1]\n\t [[{{node TopKV2}}]]"}
【发布时间】：2021-12-29 09:29:19
【问题描述】：

这是我的代码-

import tensorflow as tf    
SAVE_MODEL_PATH = "test/1"

class GetTopK(tf.Module):   
    def __init__(self):
        self.wow = "wowness!"

    @tf.function(input_signature=[
        tf.TensorSpec(shape=(None, 5), dtype=tf.float32), 
        tf.TensorSpec(shape=None, dtype=tf.int32)])
    
    def indices(self, inp_tensor, top_k):        
        indices = tf.math.top_k(input=inp_tensor, k=top_k).indices
        return {"indices": indices}


to_export = GetTopK()
tf.saved_model.save(
    to_export, 
    SAVE_MODEL_PATH, 
    signatures={"indices": to_export.indices})

它应该接受一个形状为 (None, 5) 的张量，然后使用 tf.math.top_k 函数返回具有最高值的索引

我正在使用 tensorflow-gpu==2.3.0

然后我尝试使用 docker 运行 tf 服务-

export ModelPath="$(pwd)/test" # pwd is where the above mentioned python script exists
docker run -t --rm -p 8501:8501 -v "$ModelPath:/models/nsp" -e MODEL_NAME=nsp tensorflow/serving
    
curl -d '{"signature_name": "add", "instances":[{"x": 1.5, "y": 10}]}' -X POST http://localhost:8501/v1/models/nsp:predict

它会抛出以下错误-

{
    "error": "k must be scalar, got shape [1]\n\t [[{{node TopKV2}}]]"
}

【问题讨论】：

标签： python-3.x docker tensorflow tensorflow2.0 tensorflow-serving

【解决方案1】：

我只是变了

indices = tf.math.top_k(input=inp_tensor, k=top_k).indices

到

indices = tf.math.top_k(input=inp_tensor, k=tf.squeeze(top_k)).indices

它奏效了。问题是，即使我发送一个整数，tf serving 也会将其转换为张量，所以我必须使用 tf.squeeze 将标量挤出张量

【讨论】：