How to return RaggedTensors in TF model signature and TFServing?

Viewed 383

I am using TFServing to allow prediction on my model trained with TF2. After the training I save my model with a custom signature. Let's say my input is a tensor with shape [2] and in my signature I want to add some some attributes to both elements in my input tensor:

@tf.function
def serve(input):

   // model inference

   return {
      'input': input,
      'attr': tf.constant(['attr1', 'attr2'], ['attr3', 'attr4'])
   }
signatures = {'signature': serve.get_concrete_function(
                                        tf.TensorSpec(
                                           shape=[None, 1],
                                           dtype=tf.string,
                                           name='input'))}
}
model.save(serving_dir, save_format='tf', signatures=signatures)

Using the signature above. I get the expected result, when interfacing the TFServing REST API with a single example.

{
   "predictions": [
      {
         "input": "Example Input 1",
         "attr": ["attr1", "attr2"]
      },
      {
         "input": "Example Input 2",
         "attr": ["attr3", "attr4"]
      }
   ]
}

Now, what if I want each element in my input tensor to have a variable number of attributes?

{
   "predictions": [
      {
         "input": "Example Input 1",
         "attr": ["attr1", "attr2"]
      },
      {
         "input": "Example Input 2",
         "attr": ["attr3"]
      }
   ]
}
@tf.function
def serve(input):
   return {
      'input': input,
      'attributes': tf.constant([['attr1', 'attr2'], ['attr3']])
   }

Unsurprising error when saving model:

Argument must be a tensor: [['attr1', 'attr2'], ['attr3']] - got shape [2], but wanted [2, 2]. 
@tf.function
def serve(input):
   return {
      'input': input,
      'attributes': tf.ragged.constant([['attr1', 'attr2'], ['attr3']])
   }

Error when saving:

ValueError: Got a dictionary containing non-Tensor value tf.RaggedTensor(values=Tensor('StatefulPartitionedCall:0", shape=(None,), dtype=string), row_splits=Tensor("StatefulPartitionCall:1", shape=(3,), dtype=int64)) for key date in the output of the function __inferenece_serve_26051 used to generate a SavedModel signature. Dictionaries outputs for functions used as signatures should have one Tensor output per string key.

What are my options to return variable length tensors for different inputs?

0 Answers
Related