在TensorFlow中从CSV分类特征列创建多热SparseTensor的实现方案

架构老林

2026-5-22

处理TensorFlow数据流水线中的多值稀疏ID特征转多热SparseTensor

嘿，我完全懂你在推荐系统里处理稀疏ID特征时的困扰——尤其是碰到那种用分隔符拼接的多值ID（比如示例里的aa|cc|ff），要把它们转换成多热SparseTensor，还得无缝适配TensorFlow的数据流管道对吧？我之前也踩过类似的坑，给你分享一个可行的实现方案，应该能解决你的问题。

首先明确需求：我们要把CSV中用|分隔的多值ID，映射到预设的词汇表（比如aa、bb、cc、dd、ee、ff），然后生成对应的多热SparseTensor（节省内存，适合大规模稀疏数据）。

完整实现步骤

1. 定义词汇表与映射表

首先我们需要把所有可能的ID映射成索引，用TensorFlow的lookup表来做高效的词汇转换：

import tensorflow as tf

# 预设的ID词汇表，根据你的实际业务调整
vocab = ['aa', 'bb', 'cc', 'dd', 'ee', 'ff']

# 创建静态词汇表，把ID映射成对应的整数索引
vocab_table = tf.lookup.StaticVocabularyTable(
    tf.lookup.KeyValueTensorInitializer(
        keys=vocab,
        values=tf.range(len(vocab), dtype=tf.int64)
    ),
    num_oov_buckets=1  # 预留一个桶处理不在词汇表中的未知ID
)

2. 编写CSV行解析与特征转换函数

接下来写一个处理单条CSV行的函数，完成分割多值ID、转换索引、构建SparseTensor的流程：

def parse_and_transform(line):
    # 解析CSV行，这里假设列格式是：[样本ID, 多值特征, 标签]
    sample_id, raw_features, label = tf.io.decode_csv(
        line,
        record_defaults=['', '', 0],  # 缺失值默认值
        field_delim=','
    )
    
    # 把用|分隔的多值特征拆分成单个ID
    split_ids = tf.strings.split(raw_features, sep='|')
    
    # 将ID字符串转换成词汇表中的整数索引
    feature_indices = vocab_table.lookup(split_ids)
    
    # 构建多热SparseTensor：
    # indices格式是[[样本行索引, 特征索引], ...]，单样本时行索引都是0
    indices = tf.stack([tf.zeros_like(feature_indices, dtype=tf.int64), feature_indices], axis=1)
    # values全为1，表示对应特征存在
    values = tf.ones_like(feature_indices, dtype=tf.int32)
    # dense_shape是[样本数, 词汇表大小]，单样本时是[1, 6]
    sparse_multi_hot = tf.sparse.SparseTensor(
        indices=indices,
        values=values,
        dense_shape=[1, len(vocab)]
    )
    
    # 返回转换后的特征和标签，样本ID可选返回
    return sparse_multi_hot, label

3. 构建TensorFlow数据流水线

把示例CSV数据接入流水线，测试转换效果：

# 你的示例CSV数据
csv_samples = [
    '1221,cc,1',
    '213,aa|cc|ff,1',
]

# 构建数据集并应用转换函数
dataset = tf.data.Dataset.from_tensor_slices(csv_samples)
dataset = dataset.map(parse_and_transform)

# 遍历数据集查看结果
for sparse_tensor, label in dataset:
    print(f"标签: {label.numpy()}")
    print(f"SparseTensor索引: {sparse_tensor.indices.numpy()}")
    print(f"SparseTensor值: {sparse_tensor.values.numpy()}")
    print(f"对应的稠密多热矩阵:\n{tf.sparse.to_dense(sparse_tensor).numpy()}\n")

运行结果说明

执行上面的代码后，你会得到如下输出：

标签: 1
SparseTensor索引: [[0 2]]
SparseTensor值: [1]
对应的稠密多热矩阵:
[[0 0 1 0 0 0]]
标签: 1
SparseTensor索引: [[0 0]
[0 2]
[0 5]]
SparseTensor值: [1 1 1]
对应的稠密多热矩阵:
[[1 0 1 0 0 1]]

完全符合你想要的多热编码效果！而且整个流程是纯TensorFlow操作，可以无缝集成到大规模数据的流水线中（比如从文件读取CSV、批量处理、预取等）。

批量处理注意事项

如果要处理批量数据，只需要在数据集后加上batch()操作，TensorFlow会自动合并每个样本的SparseTensor：

dataset = dataset.batch(2)
for batch_sparse, batch_labels in dataset:
    print(f"批量SparseTensor形状: {batch_sparse.dense_shape}")
    print(f"批量稠密矩阵:\n{tf.sparse.to_dense(batch_sparse).numpy()}")

这样就能得到批量的多热SparseTensor，非常适合后续的推荐模型输入（比如Wide&Deep、FM等）。

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