在 2D 张量的值高于或等于 tensorflow2 中的阈值的索引处从数组中获取文本

get texts from array at indexes where 2D tensor's values are above or equal to thresholds in tensorflow2

我有 tensorflow 张量,概率如下:

>>> valid_4_preds
array([[0.9817431 , 0.01259811, 0.50729334, 0.00053732, 0.6966804 ,
        0.00488825],
       [0.9851129 , 0.01246135, 0.38177294, 0.00378728, 0.8398497 ,
        0.68413687],
       [0.00061161, 0.00005008, 0.00017785, 0.0000152 , 0.00017121,
        0.00002404],
       [0.9991425 , 0.23962161, 0.98579687, 0.01727398, 0.9354003 ,
        0.3325037 ]], dtype=float32)

我现在需要将上述具有不同阈值的概率映射到 classes(或带有文本的张量)并获取它们。

>>> # printing classes
>>> classes
<tf.Tensor: shape=(6,), dtype=string, numpy=
array([b'class_1', b'class_2', b'class_3', b'class_4', b'class_5',
       b'class_6'], dtype=object)>
>>> # converting to bools
>>> true_falses = tf.math.greater_equal(valid_4_preds, tf.constant([0.5, 0.40, 0.20, 0.80, 0.5, 0.4]))
>>> true_falses
<tf.Tensor: shape=(4, 6), dtype=bool, numpy=
array([[ True, False,  True, False,  True, False],
       [ True, False,  True, False,  True,  True],
       [False, False, False, False, False, False],
       [ True, False,  True, False,  True, False]])>

现在,我正在尝试获取 true_falses 具有 Trues 的索引处的文本(这是我的预期输出),如下所示:

>>> <some-tensorflow-operations>
<tf.Tensor: shape=(4, 6), dtype=bool, numpy=
array([['class_1', 'class_3', 'class_5'],
       ['class_1', 'class_3', 'class_5', 'class_6'],
       [],
       ['class_1', 'class_3', 'class_5']])>

这是我尝试过的方法:

  1. tf.boolean_mask 似乎解决了目的,但它接收的 mask 严格必须是一维数组。
  2. tf.where 可用于获取索引,其输出在重塑为单一维度后可传递给 tf.gather 以获取相应的 类,如下所示:
>>> tf.gather(classes, tf.reshape(tf.where(true_falses[0] == True), shape=(-1,)))
<tf.Tensor: shape=(3,), dtype=string, numpy=array([b'class_1', b'class_3', b'class_5'], dtype=object)>

但是,我一直无法弄清楚如何在二维数组上执行此操作。

此逻辑将通过 tensorflow-serving 进入服务签名,因此操作严格只需要是 tensorflow。我如何在二维张量或数组上执行此操作?更高效、更快速的操作将不胜感激。

tf.ragged.boolean_mask?

import tensorflow as tf

classes = tf.constant([b'class_1', b'class_2', b'class_3', b'class_4', b'class_5', b'class_6'])
true_falses = tf.constant([
    [ True, False,  True, False,  True, False],
    [ True, False,  True, False,  True,  True],
    [False, False, False, False, False, False],
    [ True, False,  True, False,  True, False]]
)

tf.ragged.boolean_mask(
    data=tf.tile(tf.expand_dims(classes, 0), [tf.shape(true_falses)[0], 1]),
    mask=true_falses
)
# <tf.RaggedTensor [[b'class_1', b'class_3', b'class_5'], [b'class_1', b'class_3', b'class_5', b'class_6'], [], [b'class_1', b'class_3', b'class_5']]>