TensorFlow Datasets项目：如何加载外部TFRecord格式数据

TensorFlow Datasets项目：如何加载外部TFRecord格式数据

datasets TFDS is a collection of datasets ready to use with TensorFlow, Jax, ... 项目地址: https://gitcode.com/gh_mirrors/dat/datasets

TensorFlow Datasets（TFDS）是一个强大的工具，可以帮助开发者轻松加载和管理各种机器学习数据集。本文将详细介绍如何使用TFDS加载外部的TFRecord格式数据，这些数据可能由第三方工具生成，而非TFDS原生支持的数据集。

为什么需要加载外部TFRecord数据

在实际项目中，我们经常会遇到以下几种情况：

1. 团队内部有自定义的数据处理流水线，已经生成了TFRecord格式的数据
2. 需要使用第三方工具处理后的数据结果
3. 希望利用TFDS的统一接口来管理所有数据集，包括自定义数据

TFDS提供了完善的机制来支持这些场景，让我们能够在不修改原始数据的情况下，享受TFDS提供的各种便利功能。

准备工作

要使用TFDS加载外部TFRecord数据，需要满足以下基本条件：

1. 数据格式必须是tf.train.Example协议缓冲区格式（不支持tf.train.SequenceExample）
2. 能够用tfds.features描述数据结构
3. 按照TFDS的命名规范组织文件

文件命名规范

TFDS支持灵活的文件命名模板，通过tfds.core.ShardedFileTemplate定义。模板支持以下变量：

• {DATASET}：数据集名称
• {SPLIT}：数据分割名称（如train、test）
• {FILEFORMAT}：文件格式（如tfrecord）
• {SHARD_INDEX}：分片索引
• {NUM_SHARDS}：总分片数
• {SHARD_X_OF_Y}：分片标识（如00001-of-00005）

默认的命名模板是：{DATASET}-{SPLIT}.{FILEFORMAT}-{SHARD_X_OF_Y}

例如，MNIST数据集的TFRecord文件可能命名为：

• mnist-test.tfrecord-00000-of-00001
• mnist-train.tfrecord-00000-of-00001

添加元数据

要让TFDS正确解析TFRecord数据，我们需要提供完整的元数据信息。

定义特征结构

首先需要定义数据的特征结构，这与TFDS原生数据集的定义方式一致：

features = tfds.features.FeaturesDict({
    'image': tfds.features.Image(shape=(256, 256, 3)),
    'label': tfds.features.ClassLabel(names=['dog', 'cat']),
    'objects': tfds.features.Sequence({
        'camera/K': tfds.features.Tensor(shape=(3,), dtype=tf.float32),
    }),
})

这个结构对应于以下TFRecord特征定义：

{
    'image': tf.io.FixedLenFeature(shape=(), dtype=tf.string),
    'label': tf.io.FixedLenFeature(shape=(), dtype=tf.int64),
    'objects/camera/K': tf.io.FixedLenSequenceFeature(shape=(3,), dtype=tf.int64),
}

如果你控制数据生成流程

如果你能控制数据生成过程，可以使用serialize_example方法确保数据兼容性：

with tf.io.TFRecordWriter('path/to/file.tfrecord') as writer:
    for ex in all_exs:
        ex_bytes = features.serialize_example(data)
        writer.write(ex_bytes)

如果不控制数据生成流程

如果你无法修改数据生成过程，可以通过以下方式检查特征匹配：

# 获取人类可读的特征结构
print(features.get_serialized_info())

# 获取tf.io.parse_single_example使用的特征定义
spec = features.tf_example_spec

获取分割统计信息

TFDS需要知道每个分片中的样本数量，以支持数据集长度查询和子分割等功能。

如果你已经知道这些信息，可以直接创建SplitInfo列表：

split_infos = [
    tfds.core.SplitInfo(
        name='train',
        shard_lengths=[1024, 512],  # 各分片的样本数
        num_bytes=0,  # 数据集总大小（未知可设为0）
    ),
    # 其他分割...
]

如果不知道这些信息，可以使用compute_split_info工具自动计算。

写入元数据文件

使用write_metadata函数一次性生成所有需要的元数据文件：

tfds.folder_dataset.write_metadata(
    data_dir='/path/to/dataset/1.0.0/',
    features=features,
    split_infos=split_infos,  # 或split_infos_dir
    filename_template='{SPLIT}-{SHARD_X_OF_Y}.{FILEFORMAT}',
    description="数据集描述",
    homepage='http://example.com',
    supervised_keys=('image', 'label'),
    citation="BibTex引用格式",
)

加载数据集

从单个目录加载

生成元数据后，可以使用builder_from_directory加载数据集：

builder = tfds.builder_from_directory('/path/to/dataset/1.0.0/')

# 访问元数据
print(builder.info.splits['train'].num_examples)

# 构建数据管道
ds = builder.as_dataset(split='train[75%:]')
for ex in ds:
    # 处理数据...

从多个目录加载

对于分布在多个目录的数据（如不同代理生成的数据），可以使用builder_from_directories：

builder = tfds.builder_from_directories([
    '/path/agent1/1.0.0/',
    '/path/agent2/1.0.0/',
    '/path/agent3/1.0.0/',
])

# 使用方式与单个目录相同

推荐目录结构

为了更好的兼容性，建议采用以下目录结构：

data_dir/
    dataset_name/
        1.0.0/
        1.0.1/
    another_dataset/
        config_a/
            2.0.0/
        config_b/
            2.0.0/

这样可以直接使用tfds.load加载：

ds1 = tfds.load('dataset_name', data_dir='data_dir/')
ds2 = tfds.load('another_dataset/config_a', data_dir='data_dir/')

总结

通过本文介绍的方法，你可以轻松地将外部生成的TFRecord数据集成到TFDS生态系统中，享受统一的数据加载、分割管理和特征处理等便利功能。这种方法特别适合已有数据处理流水线但希望利用TFDS优势的项目。

datasets TFDS is a collection of datasets ready to use with TensorFlow, Jax, ... 项目地址: https://gitcode.com/gh_mirrors/dat/datasets