在TensorFlow中获取矩阵的对角线

Question

有没有办法在TensorFlow中提取方阵的对角线？也就是说,对于这样的矩阵:

[
 [0, 1, 2],
 [3, 4, 5],
 [6, 7, 8]
]

我想获取元素: [0, 4, 8]

在Numpy中,通过np.diag非常简单:

在TensorFlow中,有一个diag函数,但它只与对角线上的参数中指定的元素形成一个新的矩阵,这不是我想要的.

我可以想象如何通过跨步来实现这一目标...但我不认为在TensorFlow中争取张量.

Answer 1

使用tensorflow 0.8可以提取对角线元素tf.diag_part()(参见文档)

UPDATE

for tensorflow> = r1.12 its tf.linalg.tensor_diag_part(参见文档)

Answer 2

目前可以使用tf.diag_part提取对角线元素。这是他们的例子：

"""
'input' is [[1, 0, 0, 0],
            [0, 2, 0, 0],
            [0, 0, 3, 0],
            [0, 0, 0, 4]]
"""

tf.diag_part(input) ==> [1, 2, 3, 4]

旧答案（当 diag_part 时）不可用（如果您想实现现在不可用的东西，仍然相关）：

在查看了数学运算和张量变换之后，看起来并不存在这样的运算。即使您可以通过矩阵乘法提取此数据，它也不会有效（得到对角线是O(n)）。

您可以采用三种方法，从易到难。

1. 评估张量，使用 numpy 提取对角线，使用 TF 构建变量
2. 按照 Anurag 建议的方式使用tf.pack （也使用以下命令提取值 3tf.shape
3. 用 C++编写您自己的操作，重建 TF 并在本机使用它。