在损失函数中索引 tf 变量

Question

我在 Tensorflow 1.9.0 中定义了一个自定义损失函数（由于项目限制无法升级）。我有以下变量，经过特征值分解后获得：

# eigw.shape = (?, x)
# eigv.shape = (?, x, y)

现在，我想计算eigw 的argmax，这样

amax = tf.argmax(eigw, axis=1, output_type=tf.int32)
# amax.shape = (?,)

我想用amax 中给出的值来索引eigv，这样

# result.shape = (?, y)

我如何做到这一点？我尝试直接访问它，但这样做我遇到了形状不相等的问题。另外，我尝试使用tf.while_loop，但我是 tf 的新手，因此我没有成功。

我还有什么其他选择？如何最轻松地解决这个问题？

谢谢

Answer 1

在您的具体情况下，您可以使用任何 TensorFlow 函数来收集沿轴而不是索引的最大值。

max_value = tf.math.reduce_max(eigw, axis=1)

您可能会在文档中看到任何其他参数。由于 tesnorlfow.org 上没有更多的 TF 1.9 文档，我可以找到仍然使用静态图的 r1.15。 https://www.tensorflow.org/versions/r1.15/api_docs/python/tf/math/reduce_max

Answer 2

解决方案比预期的要简单。在深入研究了我发现的文档后，特征向量已经“按非递减顺序排序”。 所以我只需要取最后一个特征向量。感谢大家的贡献。