有没有办法在 TensorFlow 或 Keras 中提取 (k,n,n) 张量的 k 对角线？

Question

在Get the diagonal of a matrix in TensorFlow 的帖子中，对于方阵 (n,n)，一个建议是使用函数 tf.diag_part(tensor)。但是如果张量是维度 (k,n,n) ？有没有办法做到这一点？必要的输出是张量的k个方阵(n,n)的k条对角线，也就是说，我需要一个维度为(k,n)的输出。有什么建议吗？

Answer 1

您可以使用tf.map_fn，来自documentation：

映射从维度 0 的元素解包的张量列表。

所以你只需要映射tf.diag_part:

a = tf.placeholder(shape=[100, 10, 10], dtype=tf.float32)
diags = tf.map_fn(tf.diag_part, a, parallel_iterations=100)

diags的形状是(100, 10)

注意：在您的情况下，parallel_iterations 理想情况下应等于 k，以获得最佳性能。