如何計算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]