Question

是否可以在Tensorflow中使用SparseTensors进行批量矩阵乘法？我知道没有明确的函数可以执行此操作，但是我想知道是否可以使用tf.map_fn来完成此操作。

下面是我如何显式进行循环的示例：

import tensorflow as tf

adj_split = [tf.sparse.SparseTensor(indices=[[0, 1], [1, 6]], values=[1.0, 2.0], dense_shape=[10, 10]) for _ in range(5)]
adj_dense = tf.stack([tf.sparse.to_dense(adj_sp) for adj_sp in adj_split])
mu = tf.random.uniform(shape=[5, 10, 7])
mu_split = tf.unstack(mu, axis = 0)

original_result = tf.einsum('ivu, iuk->ivk', adj_dense, mu) # Batch Matrix Multiply

result_list = []
for i in range(len(adj_split)):
  result_list.append(tf.sparse.sparse_dense_matmul(adj_split[i], mu_split[i]) )

for_loop_result = tf.stack(result_list, axis = 0)

original_result == for_loop_result # True

但是，我不确定Tensorflow是否会自动并行化这些独立的for循环-如果没有，是否可以使用tf.map_fn做到这一点？我在应用tf.map_fn时遇到了麻烦，因为它似乎不允许以下形式的操作：[a1，a2，a3]，[b1，b2，b3]-> [f（a1，b1），f（a2 ，b2），f（a3，b3）]。

Tensorflow中具有稀疏张量的批矩阵乘法？

0 个答案: