Есть ли способ извлечь k диагоналей тензора (k, n, n) в TensorFlow или Keras?

В сообщении Получить диагональ матрицы в TensorFlow для квадратная матрица (n,n), одно предложение использовать функцию tf.diag_part(tensor). Но если тензор имеет размеры (k,n,n)? Существует ли способ сделать это? Необходимый вывод - это k диагоналей k квадратной матрицы (n, n) тензора, то есть мне нужен один вывод размерности (k, n). Какое-то предложение?


person OmarVP    schedule 14.04.2018    source источник


Ответы (1)


Вы можете использовать tf.map_fn из документации:

Карта на список тензоров, распакованных из elems по размерности 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 для максимальной производительности.

person squadrick    schedule 14.04.2018