python如何读取tfrecord_Tensorflow(一) TFRecord生成与读取

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