tensorflow - Tensorflow TF 中稀疏张量的 einsum

Question

我想将两个张量相乘，一个稀疏，另一个密集。稀疏的是 3D，密集的是 2D。我无法将稀疏张量转换为密集张量（即，避免使用tf.sparse.to_dense(...)）。

我的乘法由以下定律给出：

C[i,j] = \sum_k A[i,k]*B[i,k,j]

其中 C = A*B 和 A 和 B 是上面描述的密集和稀疏张量。

TF 中的执行示例如下：

# Example
# inputs
A = tf.constant([[1, 2, 3],
                 [4, 5, 6],
                 [7, 8, 9]], tf.float32)
B = tf.sparse.SparseTensor(values=tf.constant([1,1,1,1,2,1,1,1,1,1,-1], tf.float32),indices=[[0,0,1],[0,1,2],[0,2,0],[1,0,0],[1,1,1],[1,1,2],[1,2,2],[2,0,2],[2,1,1],[2,2,1],[2,2,2]], dense_shape=[3,3,3])


# output
C = tf.constant([[3, 1, 2],
                 [4, 10, 11],
                 [9, 8, -1]], tf.float32)

tf.einsum不支持稀疏张量。

我有一个版本，其中我将 3D 稀疏张量 B 切片为 2D 稀疏矩阵的集合，B[0,:,:], B[1,:,:],B[2,:,:],...并将密集矩阵 A 的每一行相乘A[i,:]，每个 2D 切片稀疏矩阵都B[i,:,:]应用该tf.sparse.sparse_dense_matmul(A[i,:],B[i,:,:])函数（切片后相应的重塑为有 2D 张量作为参数tf.sparse.sparse_dense_matmul）。然后，我堆叠所有向量结果以组装 C 矩阵。这个过程很慢并且破坏了 B 的张量结构。我想通过仅应用 Tensorflow 函数来执行相同的操作（避免循环切片和破坏稀疏张量，以便稍后通过堆叠重新组合结果）。然后，这应该与 Keras 作为神经网络的一个层一起使用（[A,B] 是输入的批处理列表，C = A*B 是该层的批处理输出）。打破张量来计算乘法对于编译图中的训练来说是疯狂的！

有任何想法吗？tf.sparse.einsum稀疏张量是否存在类似函数？

如果我将 B 转换为密集张量，应用tf.einsum(A,B,'ik,ikj->ij'). 但是，我不能失去 B 的稀疏性。

谢谢你。问候，

标签： tensorflowsparse-matrixtensorflow2.0multiplicationnumpy-einsum