如何在tf.data.Dataset中映射指定整数Tensor为新标签？

要实现将标签张量15→0、76→1、34→2的映射，需使用TensorFlow原生张量操作（图模式下无法直接解析普通Python逻辑），推荐用tf.lookup.StaticHashTable高效完成映射：

import tensorflow as tf

def relabel(label: tf.Tensor) -> tf.Tensor:
    # 定义原始标签与目标标签的映射关系
    keys = tf.constant([15, 76, 34], dtype=tf.int32)
    values = tf.constant([0, 1, 2], dtype=tf.int32)
    
    # 创建静态哈希表，指定未匹配标签的默认值
    hash_table = tf.lookup.StaticHashTable(
        tf.lookup.KeyValueTensorInitializer(keys, values),
        default_value=-1  # 可按需修改，比如设为特定标识或抛出错误
    )
    
    # 查找并返回映射后的标签
    return hash_table.lookup(label)

# 修正dataset.map的写法（原lambda语法有误）
dataset: tf.data.Dataset
dataset = dataset.map(lambda x, y: (x, relabel(y)))

关键说明： #

• tf.lookup.StaticHashTable是TensorFlow专门针对张量映射的工具，适配图模式与分布式训练场景，性能比逐条件判断更优。
• 若需校验标签范围，可添加tf.debugging.assert_in确保输入标签在指定集合内，避免出现默认值的情况。

替代方案（适合少量标签）： #

如果标签数量极少，也可以用tf.case逐条件判断：

def relabel(label: tf.Tensor) -> tf.Tensor:
    return tf.case([
        (tf.equal(label, 15), lambda: tf.constant(0)),
        (tf.equal(label, 76), lambda: tf.constant(1)),
        (tf.equal(label, 34), lambda: tf.constant(2))
    ], default=lambda: tf.constant(-1))

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