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

Question

我有两个矩阵A和B形状(M, N)非常大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和获得相同的：

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]