We got this 3D input_tensor which is a tensor representing (batch_size, N, 2).
- Where,
batch_size = total batchesN = total predictions,2 = (label, score)
I want to add the score values (2nd column elements) where labels (1st column elements) are same per batch. For example, given this tensor with 3 batches, 4 predictions per batch and 2 elements; I want required_output_tensor as result.
Condition: No for loops or tf.map_fn() for this answer. Reason, tf.map_fn() is SLOW on GPU with TF2.X. You can take a look at my sample code here that is working on 2d tensor and I can use the same with tf.map_fn().
input_tensor = tf.constant([
[
[2., 0.7],
[1., 0.1],
[3., 0.4],
[2., 0.8],
],
[
[2., 0.7],
[1., 0.1],
[1., 0.4],
[4., 0.8],
],
[
[3., 0.7],
[1., 0.1],
[3., 0.4],
[4., 0.8],
]
])
required_output_tensor = [
[
[2., 1.5],
[1., 0.1],
[3., 0.4],
],
[
[2., 0.7],
[1., 0.5],
[4., 0.8],
],
[
[3., 1.1],
[1., 0.1],
[4., 0.8],
]
]
EDIT: I can see how we will end up with ragged tensor. In that case, I'm fine with choosing top-k elements per batch where k=min(size(smallest_batch)), or it can be hard coded to topk=2.
EDIT 2: Adding additional input to try out the proposed solution:
additional_input_tensor = tf.constant([
[
[2., 0.5],
[1., 0.1],
[3., 0.4],
[2., 0.5],
],
[
[22., 0.7],
[11., 0.2],
[11., 0.3],
[44., 0.8],
],
[
[3333., 0.7],
[1111., 0.1],
[4444., 0.4],
[5555., 0.8],
],
[
[2., 0.9],
[1., 0.2],
[5., 0.3],
[2., 0.9],
]
])