TensorFlow中按索引与值转换稀疏张量的技术实现请求

问题场景 #

你手头有这样一个稀疏张量：

SparseTensorValue(indices=array([[0, 0], [1, 0], [1, 1], [1, 2]]), values=array([2, 0, 2, 5]), dense_shape=array([2, 3]))

它对应的2×3密集矩阵是：
| 2 na na |
| 0 2 5 |

现在需要生成一个2×6的新张量（因为值的全集是[0,1,2,3,4,5]共6个元素），规则是把原张量中的每个值作为新张量的列索引，对应位置设为1。最终的密集形式应该是：
| 0 0 1 0 0 0 |
| 1 0 1 0 0 1 |

解决方案 #

我们可以直接基于原稀疏张量的结构构造新的稀疏张量，核心是把原张量的values作为新张量的列索引，结合原有的行索引生成新的indices，同时把新张量的values全部设为1。

下面是用TensorFlow实现的代码：

import tensorflow as tf

# 定义原稀疏张量
original_sparse = tf.SparseTensor(
    indices=[[0, 0], [1, 0], [1, 1], [1, 2]],
    values=[2, 0, 2, 5],
    dense_shape=[2, 3]
)

# 构造新的indices：原行索引 + 原values作为新列索引
new_indices = tf.concat(
    [original_sparse.indices[:, :1],  # 提取原张量的行部分
     tf.expand_dims(original_sparse.values, axis=1)],  # 把values转为列维度
    axis=1
)

# 新张量的values全为1
new_values = tf.ones_like(original_sparse.values)

# 新张量的密集形状：行数和原张量一致，列数对应值的全集大小（6）
new_dense_shape = [2, 6]

# 创建目标稀疏张量
target_sparse = tf.SparseTensor(indices=new_indices, values=new_values, dense_shape=new_dense_shape)

# 转换为密集张量验证结果
print(tf.sparse.to_dense(target_sparse).numpy())

代码解释 #

1. 构造新索引：原稀疏张量的indices是[行, 原列]，我们只保留行部分，然后把原values作为新的列索引，拼接后得到新的indices（格式为[行, 新列]）。
2. 设置新值：因为我们需要在对应位置标记为1，所以直接生成和原values长度一致的全1数组。
3. 定义新形状：行数保持原张量的2行，列数对应值的全集大小6。

运行代码后，输出的密集矩阵就是你需要的目标结果：

[[0 0 1 0 0 0]
 [1 0 1 0 0 1]]

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