TensorFlow中如何为不同TFRecord文件配置不同解析函数？

这个问题很常见，尤其是处理多源异构TFRecord数据的时候。下面是一套清晰的实现方案，既能让你知道每个Example来自哪个文件，又能针对不同文件应用专属的解析逻辑：

1. 为每个TFRecord文件添加源标识 #

首先，我们需要给每个文件的Example都带上它的源文件名信息，这样后续解析时才能区分来源。你可以通过为每个文件创建独立Dataset，再附加文件名的方式实现：

import tensorflow as tf

filenames = ["/var/data/file1.tfrecord", "/var/data/file2.tfrecord"]

# 为每个文件创建带源标识的Dataset列表
dataset_list = []
for file_path in filenames:
    # 读取单个TFRecord文件
    single_file_ds = tf.data.TFRecordDataset(file_path)
    # 给每个Example绑定对应的源文件名
    single_file_ds = single_file_ds.map(lambda example: (example, tf.constant(file_path)))
    dataset_list.append(single_file_ds)

# 合并所有带标识的Dataset
combined_dataset = tf.data.Dataset.concatenate(*dataset_list)

如果你喜欢更简洁的写法，可以用flat_map替代循环：

filenames_tensor = tf.constant(["/var/data/file1.tfrecord", "/var/data/file2.tfrecord"])

combined_dataset = tf.data.Dataset.from_tensor_slices(filenames_tensor).flat_map(
    lambda file_path: tf.data.TFRecordDataset(file_path).map(
        lambda example: (example, file_path)
    )
)

2. 编写带来源判断的解析函数 #

现在每个Example都附带了它的源文件名，我们可以在解析函数里根据文件名选择对应的特征解析规则：

def parse_with_source(example_proto, source_file):
    # 根据源文件选择对应的特征描述
    if tf.equal(source_file, "/var/data/file1.tfrecord"):
        # file1.tfrecord对应的特征结构
        feature_spec = {
            'user_id': tf.io.FixedLenFeature([], tf.int64),
            'score': tf.io.FixedLenFeature([], tf.float32)
        }
    elif tf.equal(source_file, "/var/data/file2.tfrecord"):
        # file2.tfrecord对应的特征结构
        feature_spec = {
            'item_name': tf.io.VarLenFeature(tf.string),
            'tags': tf.io.FixedLenFeature([5], tf.int64)
        }
    else:
        # 处理未知文件的情况
        raise ValueError(f"Unrecognized source file: {source_file.numpy().decode('utf-8')}")
    
    # 解析Example
    parsed_features = tf.io.parse_single_example(example_proto, feature_spec)
    
    # 可选：如果不需要保留源文件名，只返回parsed_features即可
    return parsed_features, source_file

注意这里用tf.equal而不是普通的Python==，因为TensorFlow的图模式要求用张量操作进行判断。

3. 应用解析函数到Dataset #

最后把这个解析函数映射到合并后的Dataset上：

final_dataset = combined_dataset.map(parse_with_source)

这样处理后，每个解析后的样本都会包含（解析后的特征，源文件名），你可以根据需求选择是否保留源文件名信息。

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