通過 tensorflow 中的向量縮放稀疏張量的列

Question

假設 A 是一個稀疏的二維二進制張量 (NxN)，B 是一個向量 (1xN)。 我想按 B 縮放 A 的列： A_{i, j} <- A_{i, j} x B_{j}我正在尋找一種有效的方法來做到這一點。 目前我嘗試了兩種方法，但它們太慢了：

1. 將 B 轉換為對角矩陣 BD 並使用tf.sparse_tensor_dense_matmul(A, BD) 。 這將返回一個密集矩陣，因此該方法效率不高。
2. 將 map_fn 用於稀疏張量tf.SparseTensor(A.indices, tf.map_fn(func, (A.values, A.indices[:, 1]), dtype=tf.float32), A.dense_shape) 。 這將返回一個稀疏張量，但似乎map_fn或SparseTensor函數之一是瓶頸。

謝謝！

Answer 1

更快的解決方案可能是在您的張量B上調用gather以便能夠直接乘以A的值。 另一種選擇是使用tf.sparse.map_values ，使用tf.gather創建更新向量。

在大小為1000x1000的 99% 空稀疏矩陣上：

>>> %timeit tf.sparse.SparseTensor(A.indices, A.values*tf.gather(B,A.indices[:,1]), A.dense_shape)
855 µs ± 86 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

map_values方法：

>>> %timeit tf.sparse.map_values(tf.multiply, A, tf.gather(B,A.indices[:,1]))
978 µs ± 73 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

與sparse_dense_matmul方法相比：

>>> %timeit tf.sparse.from_dense(tf.sparse.sparse_dense_matmul(A,tf.linalg.diag(B)))
26.7 ms ± 2.24 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)