Tensorflow 批量稀疏乘法

Question

我想将稀疏张量乘以稠密张量，但要在批处理中进行。

例如，我有一个稀疏张量，其对应的密集形状为 (20,65536,65536)，其中 20 是批量大小。我想将批次中的每个 (65536,65536) 与具有密集表示的张量形状 (20,65536) 中的相应 (65536x1) 相乘。 tf.sparse_tensor_dense_matmul只接受 2 级稀疏张量。有没有办法批量执行此操作？

如果可能，由于内存限制，我想避免将稀疏矩阵转换为密集矩阵。

Answer 1

假设 a是形状为 (20, 65536, 65536) 的稀疏张量和 b具有形状的密集张量 (20, 65536) ，您可以按如下方式执行批量稀疏-密集矩阵乘法:

y_sparse = tf.sparse.reduce_sum_sparse(a * b[:, None, :], axis=-1)

此解决方案扩展了张量的第二维 b启用隐式广播。然后，通过沿最后一个轴执行稀疏-密集乘法和稀疏求和来进行批量矩阵乘法。

如 b有一个第三维，所以它是一批矩阵，你可以单独将它们的列相乘，然后将它们连接起来:

multiplied_dims = []
for i in range (b.shape[-1]):
  multiplied_dims.append(tf.expand_dims(tf.sparse.reduce_sum(a * b[:, :, i][:, None, :], axis=-1), -1))
result = tf.concat(multiplied_dims, -1)