Tf.map_fn regular tensor to ragged tensor - ragged non max supression

General Discussion
, ,
#1

Hi I want to perform a batched nms on ragged tensor.

I have following two tensors describing boxes and scores

boxes = tf.ragged.constant(
    [[[0, 0, 0, 0], [0, 0, 0, 0]], [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]],
    dtype=tf.float32,
)
scores = tf.ragged.constant([[1, 1], [1, 1, 1]], dtype=tf.float32)

I’ve tried using simply calling combined_non_max_suppression method but unsurprisingly it didn’t work.

tf.image.combined_non_max_suppression(boxes, scores, 100, 100)

It throws

ValueError: TypeError: object of type 'RaggedTensor' has no len()

So i’ve decided i could try combining tf.map_fn with non-batched verison of nms.
I can perform nms on a single batch using following code.

i = 0
tf.image.non_max_suppression(boxes[i].to_tensor(), scores[i], 100)

The issue begins when i try to combine it with map_fn

def nms(x):
    result = tf.image.non_max_suppression(x[0].to_tensor(), x[1], 100)
    result = tf.expand_dims(result, axis=0)
    result = tf.RaggedTensor.from_tensor(result)
    
    return result

output = tf.map_fn(nms, (boxes, scores), dtype=(tf.float32, tf.float32))

It throws the following

ValueError: The two structures don't have the same nested structure.

First structure: type=tuple str=(tf.float32, tf.float32)

Second structure: type=RaggedTensor str=<tf.RaggedTensor [[0, 1]]>

More specifically: Substructure "type=tuple str=(tf.float32, tf.float32)" is a sequence, while substructure "type=RaggedTensor str=<tf.RaggedTensor [[0, 1]]>" is not

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
Cell In[178], line 8
      4     result = tf.RaggedTensor.from_tensor(result)
      6     return result
----> 8 output = tf.map_fn(nms, (boxes, scores), dtype=(tf.float32, tf.float32))
...
More specifically: Substructure "type=tuple str=(tf.float32, tf.float32)" is a sequence, while substructure "type=RaggedTensor str=<tf.RaggedTensor [[0, 1]]>" is not
Entire first structure:
(., .)
Entire second structure:
.

How could i fix it ?

#3

Hi @11dk11, Could you please let me know why you are using a ragged tensor, because for bounding boxes you will get a fixed number of co-ordinates.Thank You.

#5

To fix your issue, modify the dtype parameter in tf.map_fn to match the actual output type of your nms function, which returns a RaggedTensor. Instead of specifying dtype=(tf.float32, tf.float32), you should use fn_output_signature to define the expected output structure as a RaggedTensor. Here’s a simplified correction:

output = tf.map_fn(nms, (boxes, scores), fn_output_signature=tf.RaggedTensorSpec(shape=[None], dtype=tf.int32))

This change aligns the expected output structure of tf.map_fn with the actual RaggedTensor output from your nms function, resolving the structure mismatch error.