TensorFlow中稀疏矩阵与稠密向量相乘的低内存实现方法求助

你遇到的问题核心是TensorFlow张量的不可变性——Tensor对象不能像NumPy数组那样直接通过下标赋值；而assign方法虽然能修改变量，但切片赋值的限制也让循环思路走不通，更别说这种逐元素循环在TensorFlow里效率极低，还会生成大量中间张量占用额外内存。

其实TensorFlow早就内置了专门处理稀疏-稠密矩阵乘法的高效工具，完全不用自己造轮子：tf.sparse.sparse_dense_matmul，它针对稀疏矩阵的结构做了深度优化，只计算非零元素的乘积，内存占用和计算效率都远高于手动循环。

示例代码 #

下面是一个简单的使用示例，直接对应你的需求场景：

import tensorflow as tf

# 构造你的稀疏矩阵（SparseTensor格式）
indices = sparse_mat.indices  # 你的稀疏矩阵索引
values = sparse_mat.values    # 你的稀疏矩阵非零值
dense_shape = sparse_mat.dense_shape  # 稀疏矩阵对应的稠密形状

# 构造稠密向量
dense_vec = dense_vec  # 你的输入稠密向量

# 执行稀疏-稠密乘法
# 如果dense_vec是一维，直接传入即可；如果是二维矩阵需保证维度匹配
result = tf.sparse.sparse_dense_matmul(sparse_mat, dense_vec[:, tf.newaxis])

# 运行并获取结果
with tf.Session() as sess:
    final_result = sess.run(result)
print(final_result)

关键说明 #

1. 维度匹配规则：当稀疏矩阵形状为[M, N]时，稠密张量的形状需要是[N]（一维向量）或[N, K]（矩阵），输出结果对应[M]或[M, K]，完全符合线性代数的矩阵乘法规则。
2. 内存优势：这个函数会直接利用稀疏矩阵的非零元素信息，不为零元素分配内存，也不做无效的零值乘法，比你手动循环的方式内存占用低一个量级。
3. 性能优势：TensorFlow的计算图模式下，手动循环会生成大量冗余计算节点，而内置函数是底层优化的实现，在大矩阵场景下性能差距会非常明显。

如果你的稀疏矩阵还不是SparseTensor格式，只需要用tf.SparseTensor(indices, values, dense_shape)包装即可，这一步的开销可以忽略不计。

内容的提问来源于stack exchange，提问作者sparse mat