TensorFlow中如何基于分列词集从词向量构建one hot encoder

核心思路 #

你的场景中每一列对应独立互斥的词集合，无需构建全局统一词表映射，核心逻辑为逐列生成独立One-Hot向量，再按列顺序拼接所有向量即可，完全匹配你给出的输出规则。

具体实现 #

1. 预处理全局字典 #

首先对存储列词集合的全局字典做预处理，固定每列的词顺序、生成本列词到列内局部索引的查找表、记录每列One-Hot的长度：

必须对每列的词集合做固定排序，避免Python集合无序导致每次运行映射顺序错乱。如果需要自定义词顺序，直接替换排序后的词表为你指定顺序的列表即可。

import tensorflow as tf

# 你的全局字典
global_dict = {
    0: {1, 4},
    1: {2, 5, 7},
    2: {3, 6, 8}
}

col_lookups = []  # 存储每列的词->局部索引查找层
col_vocab_depths = []  # 存储每列One-Hot向量的长度

# 按列索引从小到大遍历，保证拼接顺序和预期一致
for col_idx in sorted(global_dict.keys()):
    word_set = global_dict[col_idx]
    # 固定词表顺序，示例中排序后正好对应[1,4]、[2,5,7]、[3,6,8]
    vocab = sorted(word_set)
    col_vocab_depths.append(len(vocab))
    # 构建列专属查找层，整数输入用IntegerLookup，字符串输入替换为StringLookup即可
    lookup_layer = tf.keras.layers.IntegerLookup(
        vocabulary=tf.constant(vocab, dtype=tf.int64),
        num_oov_indices=0,  # 无未登录词时设为0，有OOV需求可设为1，对应列深度+1
        output_mode="int"
    )
    col_lookups.append(lookup_layer)

total_output_dim = sum(col_vocab_depths)

2. 编码函数 #

实现逐列编码+拼接的逻辑，支持批量输入：

def batch_onehot_encode(batch_input):
    """
    参数:
        batch_input: 形状为(batch_size, 列数)的张量，即你输入的固定长度向量批次
    返回:
        形状为(batch_size, 总输出维度)的One-Hot张量，和预期输出格式一致
    """
    onehot_segments = []
    for col_id, lookup in enumerate(col_lookups):
        # 提取当前列所有样本的值
        col_values = batch_input[:, col_id]
        # 查找值在本列词表内的局部索引
        local_idx = lookup(col_values)
        # 生成本列对应的One-Hot向量
        col_onehot = tf.one_hot(local_idx, depth=col_vocab_depths[col_id], dtype=tf.int64)
        onehot_segments.append(col_onehot)
    # 沿特征维度拼接所有列的One-Hot结果
    return tf.concat(onehot_segments, axis=-1)

3. 效果验证 #

用你给出的示例批次测试：

# 示例输入批次
demo_batch = tf.constant([
    [1, 2, 3],
    [4, 5, 6],
    [4, 7, 8]
], dtype=tf.int64)

encode_result = batch_onehot_encode(demo_batch)
print(encode_result)

输出完全匹配预期：

tf.Tensor(
[[1 0 1 0 0 1 0 0]
 [0 1 0 1 0 0 1 0]
 [0 1 0 0 1 0 0 1]], shape=(3, 8), dtype=int64)

注意事项 #

• 如果输入存在未登录词（即词不在对应列的词集合中），将IntegerLookup/StringLookup的num_oov_indices设为1，同时给对应列的col_vocab_depths值加1即可，OOV词会被映射到该列One-Hot向量的第一位。
• 该逻辑可以直接封装为Keras自定义层，集成到TensorFlow模型中实现端到端训练、推理，没有额外性能开销。

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