1. 简介
- tfrecord是一种二进制文件，能够实现数据的快速读取，是tensorflow官方推荐的一种数据处理格式；tfrecord文件中存放的基本数据是 tf.train.Example 序列化的对象；Example是Protobuf数据标准的实现。
- 一个Example消息体中包含了很多 tf.train.Feature 属性，每一个feature是key-value的键值对；key为字符串，value的数据类型如下：
tfrecord数据类型.png
tfrecord数据格式.png
2. 生成tfrecord数据
- 生成tfrecord数据格式步骤如下：https://github.com/nlpming/DeepCTR/blob/master/examples/gen_tfrecords.py
（1）首先定义每个特征的类型（tf.train.Feature -> tf.train.Int64List, tf.train.FloatList等），生成features（一个字典）；
（2）之后根据features，生成 tf.train.Example；
（3）最后写入到tfrecord文件；依赖 tf.python_io.TFRecordWriter 方法；
import tensorflow as tf
def make_example(line, sparse_feature_name, dense_feature_name, label_name):
# 1. 定义features： key -> value
features = {feat: tf.train.Feature(int64_list=tf.train.Int64List(value=[int(line[1][feat])])) for feat in
sparse_feature_name}
features.update(
{feat: tf.train.Feature(float_list=tf.train.FloatList(value=[line[1][feat]])) for feat in dense_feature_name})
features[label_name] = tf.train.Feature(float_list=tf.train.FloatList(value=[line[1][label_name]]))
# 2. 定义tf.train.Example
return tf.train.Example(features=tf.train.Features(feature=features))
def write_tfrecord(filename, df, sparse_feature_names, dense_feature_names, label_name):
# 3. 写入tfrecord文件；
writer = tf.python_io.TFRecordWriter(filename)
for line in df.iterrows():
ex = make_example(line, sparse_feature_names, dense_feature_names, label_name)
writer.write(ex.SerializeToString())
writer.close()
# write_tfrecord('./criteo_sample.tr.tfrecords',train,sparse_features,dense_features,'label')
# write_tfrecord('./criteo_sample.te.tfrecords',test,sparse_features,dense_features,'label')
3. 读取tfrecord数据
（1）定义tfrecord文件中，每个特征对应的类型；
（2）tf.data.TFRecordDataset 方法用于读取tfrecord数据格式；
（3）tf.parse_single_example 用于处理序列化后的Example对象；
# 1. 定义tfrecord文件中，存储的每个特征的格式；
feature_description = {k: tf.FixedLenFeature(dtype=tf.int64, shape=1) for k in sparse_features}
feature_description.update(
{k: tf.FixedLenFeature(dtype=tf.float32, shape=1) for k in dense_features})
feature_description['label'] = tf.FixedLenFeature(dtype=tf.float32, shape=1)
def input_fn_tfrecord(filenames, feature_description, label=None, batch_size=256, num_epochs=1, num_parallel_calls=8,
shuffle_factor=10, prefetch_factor=1,
):
# 3. tf.parse_single_example 用于处理序列化后的Example
def _parse_examples(serial_exmp):
try:
features = tf.parse_single_example(serial_exmp, features=feature_description)
except AttributeError:
features = tf.io.parse_single_example(serial_exmp, features=feature_description)
if label is not None:
labels = features.pop(label)
return features, labels
return features
def input_fn():
# 2. tf.data.TFRecordDataset用于读取tfrecord文件
dataset = tf.data.TFRecordDataset(filenames)
dataset = dataset.map(_parse_examples, num_parallel_calls=num_parallel_calls)
if shuffle_factor > 0:
dataset = dataset.shuffle(buffer_size=batch_size * shuffle_factor)
dataset = dataset.repeat(num_epochs).batch(batch_size)
if prefetch_factor > 0:
dataset = dataset.prefetch(buffer_size=batch_size * prefetch_factor)
try:
iterator = dataset.make_one_shot_iterator()
except AttributeError:
iterator = tf.compat.v1.data.make_one_shot_iterator(dataset)
return iterator.get_next()
return input_fn
4. 更多例子
4.1 BERT中tfrecord文件处理
- 参考代码：https://github.com/nlpming/bert/blob/master/run_classifier.py
- 生成tfrecord文件 和 读取tfrecord文件完整例子：
# 1. 生成tfrecord文件；
def file_based_convert_examples_to_features(
examples, label_list, max_seq_length, tokenizer, output_file):
"""Convert a set of `InputExample`s to a TFRecord file."""
writer = tf.python_io.TFRecordWriter(output_file)
for (ex_index, example) in enumerate(examples):
if ex_index % 10000 == 0:
tf.logging.info("Writing example %d of %d" % (ex_index, len(examples)))
feature = convert_single_example(ex_index, example, label_list,
max_seq_length, tokenizer)
def create_int_feature(values):
f = tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))
return f
features = collections.OrderedDict()
features["input_ids"] = create_int_feature(feature.input_ids)
features["input_mask"] = create_int_feature(feature.input_mask)
features["segment_ids"] = create_int_feature(feature.segment_ids)
features["label_ids"] = create_int_feature([feature.label_id])
features["is_real_example"] = create_int_feature(
[int(feature.is_real_example)])
tf_example = tf.train.Example(features=tf.train.Features(feature=features))
writer.write(tf_example.SerializeToString())
writer.close()
# 2. 读取tfrecord文件
def file_based_input_fn_builder(input_file, seq_length, is_training,
drop_remainder):
"""Creates an `input_fn` closure to be passed to TPUEstimator."""
name_to_features = {
"input_ids": tf.FixedLenFeature([seq_length], tf.int64),
"input_mask": tf.FixedLenFeature([seq_length], tf.int64),
"segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
"label_ids": tf.FixedLenFeature([], tf.int64),
"is_real_example": tf.FixedLenFeature([], tf.int64),
}
def _decode_record(record, name_to_features):
"""Decodes a record to a TensorFlow example."""
example = tf.parse_single_example(record, name_to_features)
# tf.Example only supports tf.int64, but the TPU only supports tf.int32.
# So cast all int64 to int32.
for name in list(example.keys()):
t = example[name]
if t.dtype == tf.int64:
t = tf.to_int32(t)
example[name] = t
return example
def input_fn(params):
"""The actual input function."""
batch_size = params["batch_size"]
# For training, we want a lot of parallel reading and shuffling.
# For eval, we want no shuffling and parallel reading doesn't matter.
d = tf.data.TFRecordDataset(input_file)
if is_training:
d = d.repeat()
d = d.shuffle(buffer_size=100)
d = d.apply(
tf.contrib.data.map_and_batch(
lambda record: _decode_record(record, name_to_features),
batch_size=batch_size,
drop_remainder=drop_remainder))
return d
return input_fn
4.2 spark生成tfrecord文件
- spark-tfrecord参考资料：https://github.com/linkedin/spark-tfrecord
- spark生成tfrecord文件例子：
from pyspark.sql.types import *
path = "test-output.tfrecord"
fields = [StructField("id", IntegerType()), StructField("IntegerCol", IntegerType()),
StructField("LongCol", LongType()), StructField("FloatCol", FloatType()),
StructField("DoubleCol", DoubleType()), StructField("VectorCol", ArrayType(DoubleType(), True)),
StructField("StringCol", StringType())]
schema = StructType(fields)
test_rows = [[11, 1, 23, 10.0, 14.0, [1.0, 2.0], "r1"], [21, 2, 24, 12.0, 15.0, [2.0, 2.0], "r2"]]
rdd = spark.sparkContext.parallelize(test_rows)
df = spark.createDataFrame(rdd, schema)
df.write.mode("overwrite").format("tfrecord").option("recordType", "Example").save(path)
df = spark.read.format("tfrecord").option("recordType", "Example").load(path)
df.show()
4.3 tf2.0 tfrecord文件处理例子
4.4 读取gz格式的tfrecord文件
def parse_data(dataset, conf):
features = {}
for fc in outputSchema:
if fc in intFeat:
feature[fc] = tf.FixedLenFeature([], tf.int64)
elif fc == "dense":
feature[fc] = tf.FixedLenFeature([denseFea_len], tf.float32)
parsed_features = tf.parse_single_example(dataset, features)
Label = parsed_features['label']
return parsed_features, label
def train_input_fn(filenames, epoch, batch_size, parallel_numbers):
Conf = load_conf()
dataset = tf.data.TFRecordDataset(filenames, compression_type='GZIP', buffer_size=10000, num_parallel_reads=10).repeat(epoch)
dataset = dataset.apply(tf.data.experimental.map_and_batch(lambda x: parse_data(x, conf), batch_size=batch_size, num_parallel_batches=10)
dataset = dataset.prefetch(tf.contrib.data.AUTOTUNE)
Iterator = dataset.make_one_shot_iterator()
return iterator.get_next()
参考资料
- TFRecord文件处理官方教程：https://www.tensorflow.org/tutorials/load_data/tfrecord?hl=zh-cn
- TFRecord处理完整例子：https://github.com/aymericdamien/TensorFlow-Examples/blob/master/tensorflow_v2/notebooks/5_DataManagement/tfrecords.ipynb
- spark-tfrecord参考资料：https://github.com/linkedin/spark-tfrecord
- BERT tfrecord文件处理：https://github.com/nlpming/bert/blob/master/run_classifier.py
- DeepCTR tfrecord文件处理：
https://github.com/nlpming/DeepCTR/blob/master/examples/gen_tfrecords.py【生成tfrecord】
https://github.com/nlpming/DeepCTR/blob/master/deepctr/estimator/inputs.py【读取tfrecord】
