如何获取稀疏张量（Sparse Tensor）每行的有效值数量？

给定如下稀疏张量：

import tensorflow as tf

st = tf.sparse.from_dense([[1, 0, 2, 5], [3, 0, 0, 4], [0, 0, 0, 0], [1, 1, 3, 0], [1, 2, 2, 2]])
print(st)

输出为：

SparseTensor(indices=tf.Tensor(
[[0 0]
 [0 2]
 [0 3]
 [1 0]
 [1 3]
 [3 0]
 [3 1]
 [3 2]
 [4 0]
 [4 1]
 [4 2]
 [4 3]], shape=(12, 2), dtype=int64), values=tf.Tensor([1 2 5 3 4 1 1 3 1 2 2 2], shape=(12,), dtype=int32), dense_shape=tf.Tensor([5 4], shape=(2,), dtype=int64))

需要将其转换为形状为(5, 1)的张量，其中每一行代表原稀疏张量对应行的有效值数量，期望结果为[3, 2, 0, 3, 4]。

解决方案 #

可以通过稀疏张量的indices属性统计每行有效值数量，具体实现代码如下：

import tensorflow as tf

st = tf.sparse.from_dense([[1, 0, 2, 5], [3, 0, 0, 4], [0, 0, 0, 0], [1, 1, 3, 0], [1, 2, 2, 2]])

# 提取所有非零元素的行索引
row_indices = st.indices[:, 0]
# 统计每个行索引的出现次数，minlength确保覆盖所有行（包括无有效值的行）
row_counts = tf.math.bincount(row_indices, minlength=st.dense_shape[0])
# 转换为(行数, 1)的目标形状
result = tf.reshape(row_counts, (-1, 1))

print(result)

执行后输出：

tf.Tensor(
[[3]
 [2]
 [0]
 [3]
 [4]], shape=(5, 1), dtype=int32)

原理说明 #

• st.indices[:, 0]提取稀疏张量中所有非零元素所在的行索引；
• tf.math.bincount会统计每个行索引出现的次数，设置minlength为稀疏张量的总行数，保证那些没有非零元素的行（比如示例中的第3行，索引为2）会被统计为0；
• 最后用tf.reshape将一维的统计结果调整为(5,1)的形状，符合需求。

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