tensorflow2 生成TFRecord

介绍

生成 TFRecord文件 分两步实现:

  1. 定义数据结构
  2. 写入文件

参见官方文档:
https://www.tensorflow.org/tutorials/load_data/tfrecord#reading_a_tfrecord_file

生成的 TFRecord文件 大小一般为 100-200MB

定义数据结构

从上图可知,最外部是 tf.train.Example,内部是tf.train.Features,再内部就是定义数据结构,该结构有些类似 JSON ,是 key:value格式,value是tf.train.Feature

例子1

写入一条记录

feature={
  "name":tf.train.Feature(bytes_list=tf.train.BytesList(value=["李雷".encode("utf-8")])),
  "age":tf.train.Feature(int64_list=tf.train.Int64List(value=[20])),
  "height":tf.train.Feature(float_list=tf.train.FloatList(value=[1.83])),
  "parents":tf.train.Feature(bytes_list=tf.train.BytesList(value=["父亲".encode("utf-8"),"母亲".encode("utf-8")])),
  "likes":tf.train.Feature(bytes_list=tf.train.BytesList(value=["篮球".encode("utf-8"),"王者荣耀".encode("utf-8"),b"LOL"]))
}

features=tf.train.Features ( feature=feature )

example=tf.train.Example ( features=features )

注意:

  • name、age 、height 虽然只有一个值,但由于 BytesList()、Int64List()、FloatList()必须传入可迭代对象,所以是list类型

  • parents 一定有2个值,父亲、母亲(在读取时会用到)

  • likes 每个人的爱好不一样,数量不等(在读取时会用到)

Feature 封装函数

上例中
tf.train.Feature(bytes_list=tf.train.BytesList(value=["父亲".encode("utf-8"),"母亲".encode("utf-8")])) 代码较长,一般用 Feature 封装函数:

def int64_feature(value):
    """
    Wrapper for inserting int64 features into Example proto.
    """
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))


def float_feature(value):
    """
    Wrapper for inserting float features into Example proto.
    """
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(float_list=tf.train.FloatList(value=value))


def bytes_feature(value):
    """
    Wrapper for inserting bytes features into Example proto.
    """
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))

感谢:
https://github.com/balancap/SSD-Tensorflow/blob/master/datasets/dataset_utils.py

例子2

写入多条记录

feature1={
  "name":bytes_feature("李雷".encode("utf-8")),
  "age":int64_feature(20),
  "height":float_feature(1.83),
  "parents":bytes_feature(["父亲".encode("utf-8"),"母亲".encode("utf-8")]),
  "likes":bytes_feature(["篮球".encode("utf-8"),"王者荣耀".encode("utf-8"),b"LOL"])
}

feature2={
  "name":bytes_feature(b"lucy"),
  "age":int64_feature(20),
  "height":float_feature(1.68),
  "parents":bytes_feature(["父亲".encode("utf-8"),"母亲".encode("utf-8")]),
  "likes":bytes_feature(["迪奥".encode("utf-8"),b"YSL",b"lv"])
}

features=[feature1,feature2]

用到上面的 Feature 封装函数

写文件

通过tf.io.TFRecordWriter 写文件

例子1

上接例子1

with tf.io.TFRecordWriter("/Users/mym/Desktop/python/1.tfrecord") as f:
    print(example.SerializeToString())
    f.write(example.SerializeToString())

例子2

上接例子2

with tf.io.TFRecordWriter("/Users/mym/Desktop/python/2.tfrecord") as f:
    for feature in features:
        features=tf.train.Features ( feature=feature )
        example=tf.train.Example ( features=features )
        print(example.SerializeToString())
        f.write(example.SerializeToString())

tf.train.Example

Protocol消息

官网api:
https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/train/Example

属性

features ,传入tf.train.Features

方法

  • SerializeToString():将对象序列化成字符串

tf.train.Features

Protocol 消息

属性

feature ,传入tf.train.Feature

tf.train.Feature

可以接受以下三种类型(参考 .proto文件):

  • tf.train.BytesList
  • tf.train.FloatList
  • tf.train.Int64List

大多数其他类型可以强制上面3个类型之一

官方api:
https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/train/Int64List

属性

  • bytes_list ,传入tf.train.BytesList
  • float_list ,传入tf.train.FloatList
  • int64_list ,传入tf.train.Int64List

tf.train.BytesList

可强制转换以下类型:
string
byte

属性

  • value,必须传入可迭代对象

tf.train.FloatList

可强制转换以下类型:
float(float32)
double(float64)

属性

  • value,必须传入可迭代对象

tf.train.Int64List

可强制转换以下类型:
bool
enum
int32
uint32
int64
uint64

属性

  • value,必须传入可迭代对象

tf.io.TFRecordWriter

写TFRecord文件,支持with语句

官方api:
https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/io/TFRecordWriter

函数

  • __init__ (path,options=None) :初始化
  • write():传入tf.train.Example序列化后的字符串
  • flush()
  • close()

原文出处:/show_1EF57aDz46fN.html