如何在 Tensorflow 中使用另一个张量从张量中选择值

Question

我有 2 个具有相同形状的数组（张量），例如：

array1 = [[[0, 1, 0], [1, 0, 0]], [[0, 1, 0], [0, 0, 1]]]
array2 = [[[0.2, 0.7, 0.1], [0.2, 0.4, 0.4]], [[0.2, 0.8, 0.0], [0.5, 0.2, 0.3]]]

我想创建一个新数组（张量），以便根据

array2

值选择

array1

中的值。

array3 = tf.magic_function(array1, array2)
# array3 = [[[0.7], [0.2]], [[0.4], [0.3]]]

我也有 array1 不是单热编码的，因此假设

array1 = [[1, 0], [1, 2]]

的解决方案也是可以接受的。该解决方案必须使用 Tensorflow 函数，NumPy 是不可接受的，因为由于急于执行，我无法在训练期间将我的张量转换为 NumPy 数组。这是针对加权（稀疏）分类交叉熵损失函数。

我正在使用

Tensorflow 2.10

和

Python 3.9

。

向上查找

tf.gather()

和

tf.nd_gather()

，但那些似乎没有达到我想要的效果。

Answer 1

以下代码似乎可以解决问题：

import tensorflow as tf

array1 = [[[0, 1, 0], [1, 0, 0]], [[0, 1, 0], [0, 0, 1]]]
array2 = [[[0.2, 0.7, 0.1], [0.2, 0.4, 0.4]], [[0.2, 0.8, 0.0], [0.5, 0.2, 0.3]]]

tensor1 = tf.convert_to_tensor(array1, dtype=tf.float32)
tensor2 = tf.convert_to_tensor(array2)

array3 = tf.reduce_max(tf.math.multiply(tensor1, tensor2), axis=2, keepdims=True)
print(array3)
# tf.Tensor(
# [[[0.7]
#   [0.2]]
#  [[0.8]
#   [0.3]]], shape=(2, 2, 1), dtype=float32)