如何在TensorFlow中实现稀疏矩阵每行与稠密向量逐元素相乘

实现方案 #

你完全不需要将大尺寸稀疏矩阵A转为稠密矩阵，仅基于A的稀疏元数据即可完成计算，计算成本仅和A的非零元素数量正相关，不会额外占用稠密矩阵的内存开销。

核心逻辑是：A的所有非零元素都对应明确的列索引，直接用列索引从向量B中取出对应值，和A原有的非零值逐元素相乘后，复用A的索引和形状构造新的稀疏矩阵即可。

代码实现 #

import tensorflow as tf

# 你原有的A和B定义
A = tf.sparse.SparseTensor(indices=[[0,0],[1,1],[1,2]], values=[1,1,1], 
                dense_shape=[2,3])
B = tf.constant([4,3,5])

# 核心计算逻辑
# 提取A所有非零元素对应的列索引
col_idx = A.indices[:, 1]
# 取出B中对应列的数值，和A的非零值相乘得到新值
new_values = A.values * tf.gather(B, col_idx)
# 构造结果稀疏矩阵
C = tf.sparse.SparseTensor(
    indices=A.indices,
    values=new_values,
    dense_shape=A.dense_shape
)

# 可选：如果后续需要用到TensorFlow的其他稀疏算子，建议补一步重排，保证稀疏索引有序
C = tf.sparse.reorder(C)

执行后得到的C和你预期的输出完全一致。

内容的提问来源于stack exchange，提问作者Jiawei Lu