点产品占用大量内存

Question

我有两个大小为[I，J]的2D张量A和B，我想对A的每个第j个向量与B的第j个向量进行点积。这消耗了大量内存。 我正在这样做：

for j in range (J):
      K.dot(K.reshape(A[:,j], [1,-1]),K.reshape(B[:,j], [-1,1]))

有更好的方法吗？

Answer 1

您可以使用tf.einsum：

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