【数据集处理】TFRecord 文件

藓类少女

于 2024-08-13 16:25:59 发布

阅读量556

点赞数 13

分类专栏： # 数据集处理文章标签：数据集处理 tensorflow

本文链接：https://blog.csdn.net/a13545564067/article/details/140996930

版权

数据集处理专栏收录该内容

5 篇文章 0 订阅

订阅专栏

TFRecord 文件是 TensorFlow 中的一种常用数据格式，主要用于高效地存储大规模数据集。它可以将数据存储为二进制格式，既减少了存储空间，又加快了数据读取速度，尤其适合大规模的机器学习任务。

一、TFRecord 的特点

高效的存储格式：TFRecord 以二进制格式存储数据，相比于文本格式，如 CSV 或 JSON，可以更有效地利用磁盘空间。
便于与 TensorFlow 集成：TFRecord 格式与 TensorFlow 紧密集成，使用 TensorFlow 的 API 可以方便地读取和解析该格式的数据。
适合大规模数据：TFRecord 文件支持顺序读写操作，非常适合处理大规模数据集，比如深度学习任务中的图像、文本或语音数据。

二、TFRecord 文件结构

TFRecord 文件本质上是一个二进制序列化文件。它的每一条记录都是一个序列化后的 tf.train.Example，而 tf.train.Example 是一个 Protocol Buffer（Protobuf）消息，里面包含了数据的键值对。

典型的 tf.train.Example 包含如下结构：

Features: 一个键值对集合，其中键是字符串，值可以是多种类型（字符串、整数、浮点数等）。

import tensorflow as tf

# 创建一个 Example 对象
example = tf.train.Example(features=tf.train.Features(feature={
    'feature_name': tf.train.Feature(int64_list=tf.train.Int64List(value=[value])),
    # 可以添加更多的 feature
}))

三、读写 TFRecord 文件

写入数据到 TFRecord 文件

可以通过 tf.io.TFRecordWriter 将数据写入到 TFRecord 文件中。

import tensorflow as tf

with tf.io.TFRecordWriter('output.tfrecord') as writer:
    for i in range(num_samples):
        example = tf.train.Example(features=tf.train.Features(feature={
            'feature_name': tf.train.Feature(int64_list=tf.train.Int64List(value=[i]))
        }))
        writer.write(example.SerializeToString())

从 TFRecord 文件读取数据

可以使用 tf.data.TFRecordDataset 结合 tf.parse_single_example 来解析和读取 TFRecord 文件中的数据。

import tensorflow as tf

# 定义解析函数
def _parse_function(proto):
    # 定义 feature 的描述
    keys_to_features = {
        'feature_name': tf.io.FixedLenFeature([], tf.int64)
    }
    # 解析数据
    parsed_features = tf.io.parse_single_example(proto, keys_to_features)
    return parsed_features

# 创建 Dataset 对象
dataset = tf.data.TFRecordDataset('output.tfrecord')
dataset = dataset.map(_parse_function)

# 遍历数据
for record in dataset:
    print(record['feature_name'].numpy())

四、压缩 TFRecord 文件

在 TensorFlow 中，TFRecord 文件可以通过多种压缩方式进行压缩，以减少磁盘空间占用并加快读取速度。常见的压缩方式包括 GZIP 和 ZLIB。在读取或写入 TFRecord 文件时，可以指定这些压缩方式。

写入带压缩的 TFRecord 文件

当你在写入 TFRecord 文件时，可以通过指定 options 参数来使用压缩方式。

import tensorflow as tf

# 使用 GZIP 压缩
options = tf.io.TFRecordOptions(compression_type="GZIP")

with tf.io.TFRecordWriter('output.tfrecord.gz', options=options) as writer:
    for i in range(num_samples):
        example = tf.train.Example(features=tf.train.Features(feature={
            'feature_name': tf.train.Feature(int64_list=tf.train.Int64List(value=[i]))
        }))
        writer.write(example.SerializeToString())

这里，compression_type 可以是 "GZIP" 或 "ZLIB"。

读取带压缩的 TFRecord 文件

在读取带压缩的 TFRecord 文件时，也需要指定相应的压缩方式：

import tensorflow as tf

# 使用 GZIP 压缩读取
dataset = tf.data.TFRecordDataset('output.tfrecord.gz', compression_type="GZIP")

# 定义解析函数
def _parse_function(proto):
    keys_to_features = {
        'feature_name': tf.io.FixedLenFeature([], tf.int64)
    }
    parsed_features = tf.io.parse_single_example(proto, keys_to_features)
    return parsed_features

# 应用解析函数
dataset = dataset.map(_parse_function)

# 遍历数据
for record in dataset:
    print(record['feature_name'].numpy())