我正在尝试从 Tensorflow 的内置稀疏矩阵乘法 API 中获得性能优势。 并且keveman 建议 tf.embedding_lookup_sparse 是正确的方法。
但是，embedding_lookup_sparse 的性能似乎对my experiments 有点失望。虽然它执行的矩阵乘法相当小， 和 ，但稀疏度为 0.1 的稀疏 matmul 无法赢得密集矩阵乘法。
如果我的实现是正确的，我认为原因之一是 Tensorflow 使用了 COO 格式，它保存了所有索引非零对。我不是这个领域的专家，但是，CSR 格式在这种计算上的性能不是众所周知的吗？ Tensorflow内部使用CSR以外的COO格式进行稀疏矩阵表示有什么明显的原因吗？
【问题讨论】：
为了记录，您说矩阵乘法，但您的一个矩阵实际上是一个向量 (1 x 3196)。所以这将使它成为一个 矩阵向量乘法（不同的 BLAS 内核）。我假设你的意思是矩阵向量乘法作为我的答案。
是的，对于矩阵向量乘法，CSR理论上应该比 COO 快；这是因为 CSR 格式的存储大小是 O(2nnz + n) 与 O(3nnzs) 并且稀疏矩阵向量乘法在许多情况下是内存受限的。
与密集矩阵乘法相比，确切的性能差异会因问题大小、稀疏模式、数据类型和实现而异。很难直接说哪个应该更快，因为稀疏存储格式引入了间接性，这可能会导致局部性降低和算术单元利用率低（例如，不使用矢量化）。
特别是当矩阵和向量非常小以至于几乎所有东西都可以放入缓存时，我预计性能收益有限。稀疏矩阵结构通常对真正的大矩阵更有用，从 10sK x 10sK 到 1B x 1B，使用密集表示甚至不适合主内存。对于小问题，根据我的经验，与密集格式相比的存储优势通常会被局部性和算术效率的损失所抵消。在某种程度上，这是通过混合存储格式（例如 Block CSR）来解决的，它试图两全其美，并且对某些应用程序非常有用（看起来 tensorflow 不支持这一点）。
在tensorflow 中，我假设使用COO 格式，因为它对其他操作更有效，例如它支持O(1) 数据结构的更新、插入和删除。在稀疏矩阵向量乘法中以大约 50% 的性能进行交易以提高这些操作的性能似乎是合理的。
【讨论】：