如何计算 Tensorflow 中矩阵乘积的对角线？

Question

我有两个形状为 (M, N) 的矩阵 A 和 B，非常大 M 和很小 N。

我想将它们相乘然后对结果取对角线：

C = tf.matmul(A, B)
D = tf.diag_part(C)

不幸的是，这需要创建非常大的 (M, M) 矩阵，内存无法容纳。

但是我不需要这些数据中的大部分。那么，是否可以一步算出这个值呢？

有没有类似einsum但没有求和的东西？

Answer 1

你需要的相当于：

tf.einsum('ij,ij->i', A, B)

或：

tf.reduce_sum(A * B, axis=1)

示例：

A = tf.constant([[1,2],[2,3],[3,4]])
B = tf.constant([[3,4],[1,2],[2,3]])

with tf.Session() as sess:
    print(sess.run(tf.diag_part(tf.matmul(A, B, transpose_b=True)))) 
# [11  8 18]

with tf.Session() as sess:
    print(sess.run(tf.reduce_sum(A * B, axis=1)))
#[11  8 18]

with tf.Session() as sess:
    print(sess.run(tf.einsum('ij,ij->i', A, B)))
#[11  8 18]

Answer 2

你可以使用A和B transpose的dot product得到相同的：

tf.reduce_sum(tf.multiply(A, tf.transpose(B)), axis=1)

代码：

import tensorflow as tf
import numpy as np

A = tf.constant([[1,4, 3], [4, 2, 6]])
B = tf.constant([[5,4,],[8,5], [7, 3]])

E = tf.reduce_sum(tf.multiply(A, tf.transpose(B)), axis=1)

C = tf.matmul(A, B)
D = tf.diag_part(C)
sess = tf.InteractiveSession()

print(sess.run(D))
print(sess.run(E))

#Output
#[58 44]
#[58 44]

如何计算 Tensorflow 中矩阵乘积的对角线？

How to calculate diagonal of a matrix product in Tensorflow?

python

matrix

diagonal

tensorflow