TFRecord生成
一、为什么使用TFRecord?
正常情况下我们训练文件夹经常会生成 train, test 或者val文件夹,这些文件夹内部往往会存着成千上万的图片或文本等文件,这些文件被散列存着,这样不仅占用磁盘空间,并且再被一个个读取的时候会非常慢,繁琐。占用大量内存空间(有的大型数据不足以一次性加载)。此时我们TFRecord格式的文件存储形式会很合理的帮我们存储数据。TFRecord内部使用了“Protocol Buffer”二进制数据编码方案,它只占用一个内存块,只需要一次性加载一个二进制文件的方式即可,简单,快速,尤其对大型训练数据很友好。而且当我们的训练数据量比较大的时候,可以将数据分成多个TFRecord文件,来提高处理效率。
二、 生成TFRecord简单实现方式
我们可以分成两个部分来介绍如何生成TFRecord,分别是TFRecord生成器以及样本Example模块。
TFRecord生成器
writer = tf.python_io.TFRecordWriter(record_path)
writer.write(tf_example.SerializeToString())
writer.close()
这里面writer就是我们TFrecord生成器。接着我们就可以通过writer.write(tf_example.SerializeToString())来生成我们所要的tfrecord文件了。这里需要注意的是我们TFRecord生成器在写完文件后需要关闭writer.close()。这里tf_example.SerializeToString()是将Example中的map压缩为二进制文件,更好的节省空间。那么tf_example是如何生成的呢?那就是下面所要介绍的样本Example模块了。
Example模块
首先们来看一下Example协议块是什么样子的。
message Example {
Features features = 1;
};
message Features {
map feature = 1;
};
message Feature {
oneof kind {
BytesList bytes_list = 1;
FloatList float_list = 2;
Int64List int64_list = 3;
}
};
我们可以看出上面的tf_example可以写入的数据形式有三种,分别是BytesList, FloatList以及Int64List的类型。那我们如何写一个tf_example呢?下面有一个简单的例子。
def int64_feature(value):
return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
def bytes_feature(value):
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
tf_example = tf.train.Example(
features=tf.train.Features(feature={
'image/encoded': bytes_feature(encoded_jpg),
'image/format': bytes_feature('jpg'.encode()),
'image/class/label': int64_feature(label),
'image/height': int64_feature(height),
'image/width': int64_feature(width)}))
下面我们来好好从外部往内部分解来解释一下上面的内容。
(1)tf.train.Example(features = None) 这里的features是tf.train.Features类型的特征实例。
(2)tf.train.Features(feature = None) 这里的feature是以字典的形式存在,*key:要保存数据的名字 value:要保存的数据,但是格式必须符合tf.train.Feature实例要求。
三、 生成TFRecord文件完整代码实例
首先我们需要提供数据集
图片文件夹
通过图片文件夹我们可以知道这里面总共有七种分类图片,类别的名称就是每个文件夹名称,每个类别文件夹存储各自的对应类别的很多图片。下面我们通过一下代码(generate_annotation_json.py和generate_tfrecord.py)生成train.record。
generate_annotation_json.py
# -*- coding: utf-8 -*-
# @Time : 2018/11/22 22:12
# @Author : MaochengHu
# @Email : wojiaohumaocheng@gmail.com
# @File : generate_annotation_json.py
# @Software: PyCharm
import os
import json
def get_annotation_dict(input_folder_path, word2number_dict):
label_dict = {}
father_file_list = os.listdir(input_folder_path)
for father_file in father_file_list:
full_father_file = os.path.join(input_folder_path, father_file)
son_file_list = os.listdir(full_father_file)
for image_name in son_file_list:
label_dict[os.path.join(full_father_file, image_name)] = word2number_dict[father_file]
return label_dict
def save_json(label_dict, json_path):
with open(json_path, 'w') as json_path:
json.dump(label_dict, json_path)
print("label json file has been generated successfully!")
generate_tfrecord.py
# -*- coding: utf-8 -*-
# @Time : 2018/11/23 0:09
# @Author : MaochengHu
# @Email : wojiaohuma