Question

我有一个看起来像这样的张量：

arr = tf.convert_to_tensor([[3, 1],
                            [6, 2],
                            [1, 1],
                            [3, 0],
                            [5, 1],
                            [1, 0],
                            [4, 2]])

我需要仅在第二个元素等于1时计算每个单元格的第一个元素的平均值。因此，在此示例中，我将选择单元格[3, 1], [1, 1], [5, 1]并计算average(3, 1, 5)=3。然后，我需要对以0结尾然后是2结尾的单元格重复此操作，并总结所有平均值。

我将此作为神经网络中损失函数的组件编写，因此我需要仅使用TensorFlow函数执行此操作。我在Numpy中解决了这个问题，它看起来像这样：

means_sum = sum([np.mean(np.extract(arr.transpose()[1] == rank, arr.transpose()[0])) for rank in [0, 1, 2]])

但是我找不到在TensorFlow中做到这一点的方法。有什么想法吗？

Answer 1

所以基本上你想要

编写一个条件语句来匹配每个元素的第二个索引。我会尝试tf.gather和tf.not_equal（或等效函数）来过滤匹配元素，这些元素返回与（0,1）值中的输入相同形状的矩阵，然后乘以元素来过滤输入
按第一个索引平均元素（tf.reduce_mean应该足够）

你如何基于其内容划分张量？

1 个答案: