如何在TensorFlow的tf.data.Dataset.map预处理中将字符串标签映射为int类编号

解决方法 #

报错的核心原因是 tf.data.Dataset.map 默认运行在图执行模式下，无法调用仅 eager 模式可用的 .numpy() 接口。你可以使用 TensorFlow 内置的静态哈希表 tf.lookup.StaticHashTable 实现字符串到整数的映射，它完全兼容图模式，不需要任何 eager 操作，也不需要修改原有 TFRecord 的存储结构。

以下是修改后的可运行代码：

#!/usr/bin/env python3
import tensorflow as tf
import functools

# 原有标签映射逻辑不变
example_label_map = {
    "bulldog": "dog",
    "labrador": "dog",
    "persian cat": "cat",
    "cow": "neg",
}

labels_to_class_int = {
    "neg": 0,   # 负类
    "dog": 1,
    "cat": 2,
}

def parse_tfrec_function(example, lookup_table):
    image_feature_description = {
        "image": tf.io.FixedLenFeature([], tf.string),
        "label": tf.io.FixedLenFeature([], tf.string)
    }

    features = tf.io.parse_single_example(example, image_feature_description)
    image = tf.io.parse_tensor(features["image"], tf.uint8)
    
    # 替换原有逻辑：直接用lookup表查询得到类别编号
    class_num = lookup_table.lookup(features["label"])
    return image, class_num


tfrec_file = ['/path/to/dataset.tfrec']
dataset = tf.data.TFRecordDataset(tfrec_file)

# 构建标签映射字典
labels_map = {}
for k, v in example_label_map.items():
    labels_map[k] = labels_to_class_int[v]

# --------------- 新增部分：构建静态哈希表 ---------------
# 提取所有键和值转为张量
keys = tf.constant(list(labels_map.keys()))
values = tf.constant(list(labels_map.values()), dtype=tf.int64)
# 初始化键值对
init = tf.lookup.KeyValueTensorInitializer(keys, values)
# 构建静态哈希表，找不到匹配键时默认返回负类编号0
lookup_table = tf.lookup.StaticHashTable(init, default_value=0)
# --------------------------------------------------------

# 传入lookup表到解析函数
parser_fn = functools.partial(parse_tfrec_function, lookup_table=lookup_table)
parsed_dataset = dataset.map(parser_fn)

方案优势 #

1. 完全兼容图执行模式，不会触发eager相关报错，预处理效率高
2. 不需要修改原有TFRecord的存储内容，切换分类任务（比如从粗分类切换为细分类）时，只需要修改构建labels_map的逻辑，重新生成哈希表即可，不需要生成多份TFRecord文件
3. 支持自定义默认值，可处理未出现在映射表中的未知标签

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