如何将特定列具有相似值的张量元素合并?

5
我们得到了一个3D的input_tensor,它是一个表示(batch_size, N, 2)的张量。
其中, - batch_size = 总批次数 - N = 总预测数 - 2 = (标签,分数) 我想要在每个批次中,将标签相同的分数值(第二列元素)相加。例如,给定这样一个张量,有3个批次,每个批次有4个预测和2个元素;我想要得到required_output_tensor作为结果。
条件:不能使用for循环tf.map_fn(),原因是tf.map_fn()在TF2.X上GPU上运行速度较慢。您可以查看我的示例代码,该代码适用于2D张量,并且我可以在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],
    ]
]

编辑:我可以看出我们最终会得到不规则张量。在这种情况下,我认为可以选择每个批次的前k个元素,其中k=min(size(smallest_batch)),或者可以硬编码为topk=2。

编辑2:添加额外的输入以尝试提出的解决方案:

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],
    ]
])

通常情况下,这个问题并没有被很好地定义,除非你使用了不规则张量,因为每个组中重复值的数量可能不同,这将导致每个组的聚合结果数量不同。一个可能的解决方案是在所有组中都有所有标签,并用零填充缺失标签的分数。 - jdehesa
@jdehesa 是的,我能看到。我忘记包括那种情况,即我们最终得到了不规则张量。在这种情况下,我可以选择每批选择topk个元素。我已经在原始问题中添加了澄清注释。谢谢! - Snehal
1
@Snehal,你有多少个类?(4个?) - Marco Cerliani
1
@MarcoCerliani:总课程数量是任意的,很遗憾。如果我给你一个确定的数字,会有帮助吗?如果您有解决方案,可以在您的答案中使用 total_classes=50 - Snehal
2个回答

2

虽然不完全符合您的要求,但如果您知道类别数量,并且不想有不规则张量,您可以使用独热编码为相同的类别添加不同的分数:

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],
    ]
])


number_of_classes = 5

#first split the labels from scores
labels = tf.expand_dims(input_tensor[:,:,0], axis=-1)
scores = tf.expand_dims(input_tensor[:,:,1], axis=-1)

#get a one-hot encoding for the labels
#the way you do this would likely depend on your specific labels
#the way I do it here is not very robust (maybe use half open intervals instead)
class_indices = tf.reshape(tf.range(number_of_classes, dtype=tf.float32), shape=(1,1,number_of_classes))
one_hots = tf.cast(tf.equal(class_indices, labels), tf.float32)
print(one_hots.shape)  # (batch, N, number_of_classes)

#now multiply the one hots by the scores, and add all together
scored_one_hots = scores * one_hots
scores_per_index = tf.reduce_sum(scored_one_hots, axis=1) # (batch, number_of_classes) 
# where the second index denotes the class and contains the score for that class

# now finish up by combining these scores with the labels
# edit: of course this part too depends on how you actually did the encoding
batch_size = input_tensor.shape[0]
ordered_labels = tf.repeat(tf.expand_dims(tf.range(number_of_classes, dtype=tf.float32), axis=0), batch_size, axis=0)

result = tf.stack([ordered_labels, scores_per_index], axis=2)
print(result)

打印结果:

(3, 4, 5)
tf.Tensor(
[[[0.  0. ]
  [1.  0.1]
  [2.  1.5]
  [3.  0.4]
  [4.  0. ]]

 [[0.  0. ]
  [1.  0.5]
  [2.  0.7]
  [3.  0. ]
  [4.  0.8]]

 [[0.  0. ]
  [1.  0.1]
  [2.  0. ]
  [3.  1.1]
  [4.  0.8]]], shape=(3, 5, 2), dtype=float32)

你制作one-hot向量的方式取决于标签的具体情况(tf.equals可能不是最好的选择,但你可以使用比较等方法)。

虽然这是一个可以接受的解决方案,当你知道标签的编码或范围时,但在处理500K个标签或标签没有特定顺序时,它将无法很好地扩展。无论如何,感谢您的答案。这是对问题的很好的回答。 - Snehal

1
这个问题一般来说不太明确,因为输入组中非重复id值的数量可能不同,所以结果不会是一个密集张量。您可以尝试使用不规则张量,但这可能有限制。一种选项是生成每个输出组都具有每个id的结果,对于与相应输入组不匹配的id的得分将被简单设置为零。以下是如何实现此操作的方法:
import tensorflow as tf

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],
    ]
])
# Take input tensor shape
s = tf.shape(input_tensor)
# Flatten first dimensions
flat = tf.reshape(input_tensor, (-1, 2))
# Find unique id values
group_ids, group_idx = tf.unique(flat[:, 0], out_idx=s.dtype)
# Shift id indices per group in the input
num_groups = tf.reduce_max(group_idx) + 1
group_shift = tf.tile(tf.expand_dims(num_groups * tf.range(s[0]), 1), (1, s[1]))
group_idx_shift = group_idx + tf.reshape(group_shift, (-1,))
# Aggregate per group in the input
num_groups_shift = num_groups * s[0]
# Either use unsorted_segment_sum
group_sum = tf.math.unsorted_segment_sum(flat[:, 1], group_idx_shift, num_groups_shift)
# Or use bincount
group_sum = tf.math.bincount(group_idx_shift, weights=flat[:, 1],
                             minlength=num_groups_shift)
# Reshape and concatenate
group_sum_res = tf.reshape(group_sum, (s[0], num_groups))
group_ids_res = tf.tile(tf.expand_dims(group_ids, 0), (s[0], 1))
result = tf.stack([group_ids_res, group_sum_res], axis=-1)
# Sort results
result_s = tf.argsort(group_sum_res, axis=-1, direction='DESCENDING')
result_sorted = tf.gather_nd(result, tf.expand_dims(result_s, axis=-1), batch_dims=1)
print(result_sorted.numpy())
# [[[2.  1.5]
#   [3.  0.4]
#   [1.  0.1]
#   [4.  0. ]]
# 
#  [[4.  0.8]
#   [2.  0.7]
#   [1.  0.5]
#   [3.  0. ]]
# 
#  [[3.  1.1]
#   [4.  0.8]
#   [1.  0.1]
#   [2.  0. ]]]

编辑:

这里提供一种使用不规则张量输出的替代方法:

import tensorflow as tf

input_tensor = tf.constant([...])
# Same as before
s = tf.shape(input_tensor)
flat = tf.reshape(input_tensor, (-1, 2))
group_ids, group_idx = tf.unique(flat[:, 0], out_idx=s.dtype)
num_groups = tf.reduce_max(group_idx) + 1
group_shift = tf.tile(tf.expand_dims(num_groups * tf.range(s[0]), 1), (1, s[1]))
group_idx_shift = group_idx + tf.reshape(group_shift, (-1,))
# Apply unique again to find ids per batch
group_ids2_ref, group_idx2 = tf.unique(group_idx_shift)
group_ids2 = tf.gather(group_ids, group_ids2_ref % num_groups)
# Also can use unsorted_segment_sum here if preferred
group_sum = tf.math.bincount(group_idx2, weights=flat[:, 1])
# Count number of elements in each output group
out_sizes = tf.math.bincount(group_ids2_ref // num_groups, minlength=s[0])
# Make ragged result
group_sum_r = tf.RaggedTensor.from_row_lengths(group_sum, out_sizes)
group_ids_r = tf.RaggedTensor.from_row_lengths(group_ids2, out_sizes)
result = tf.stack([group_ids_r, group_sum_r], axis=-1)
print(*result.to_list(), sep='\n')
# [[2.0, 1.5], [1.0, 0.10000000149011612], [3.0, 0.4000000059604645]]
# [[2.0, 0.699999988079071], [1.0, 0.5], [4.0, 0.800000011920929]]
# [[3.0, 1.100000023841858], [1.0, 0.10000000149011612], [4.0, 0.800000011920929]]

2
@Snehal 我提供了两种选择,因为我不确定它们之间是否有任何显着的区别,这可能取决于数组的形状和类别数量。颁发奖励并没有什么急迫性,您可以保留剩余的时间来吸引其他潜在答案。 - jdehesa
1
@Snehal 我添加了一种替代实现,它会产生一个不规则张量作为结果。如果您总共有大量不同的ID,但每个批次中只有很少的ID,则这可能非常有用。 - jdehesa
似乎你提出的两个解决方案没有根据第二列对第一列进行排名。我已经添加了一个更大的“input_tensor”来测试上述两个解决方案。为了简单起见,我还附上了这个要点,其中包含了我们讨论的所有内容。它显示每批次的前3个结果,我通过第二列对最终的“result”进行排序,返回每批次的前3个元素。https://gist.github.com/spate141/4ea446a7cf0d501811ed6cc20b214226 - Snehal
1
@Snehal 抱歉,我发现我在最初的排序方面做错了什么,实际上这并不是必要的,因为你想要按分数排序结果。我已经使用你提供的新输入示例进行了测试,现在我认为它是正确的。 - jdehesa
1
太准了!现在它完美了,通过了我投掷给它的所有测试用例!如果我能给你超过50个赏金点就好了!这是送给您的链接,先生:https://rb.gy/fdgjsh - Snehal
显示剩余4条评论

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接