TensorFlow中的TFrecord

tfrecord数据文件是一种将图像数据和标签统一存储的二进制文件，能更好的利用内存，在tensorflow中快速的复制，移动，读取，存储等。tfrecord文件包含了tf.train.Example 协议缓冲区(protocol buffer，协议缓冲区包含了特征 Features)。你可以写一段代码获取你的数据，将数据填入到Example协议缓冲区(protocol buffer)，将协议缓冲区序列化为一个字符串，并且通过tf.python_io.TFRecordWriter class写入到TFRecords文件。tf.train.Example的定义如下：

message Example {
 Features features = 1;
};
 
message Features{
 map<string,Feature> featrue = 1;
};
 
message Feature{
    oneof kind{
        BytesList bytes_list = 1;
        FloatList float_list = 2;
        Int64List int64_list = 3;
    }
};

从上述代码可以看出，tf.train.Example中包含了属性名称到取值的字典，其中属性名称为字符串，属性的取值可以为字符串（BytesList）、实数列表（FloatList）或者整数列表（Int64List）。

保存数据：

首先需要给定tfrecord文件名称，并创建一个文件，之后就可以创建一个循环来依次写入数据。值得注意的是赋值给example的数据格式，从前面tf.train.Example的定义可知，tfrecord支持整型、浮点数和二进制（即字符串）三种格式。对于图片等数组形式（array）的数据，可以保存为numpy array的格式，转换为string，然后保存到二进制格式的feature中。对于单个的数值（scalar），可以直接赋值。这里value=[×]的[]非常重要，也就是说输入的必须是列表（list）。当然，对于输入数据是向量形式的，可以根据数据类型（float还是int）分别保存。并且在保存的时候还可以指定数据的维数。

读取数据：

从TFRecords文件中读取数据，首先需要用tf.train.string_input_producer生成一个解析队列。之后调用tf.TFRecordReader的read()方法首先读取解析队列，返回serialized_example对象，之后调用tf.parse_single_example操作将Example协议缓冲区(protocol buffer)解析为张量。然后调用 tf.decode_raw解码，将二进制转为整形，然后改变形状即可得到像素矩阵。

保存和读取TFrecord文件的完整代码如下：

#coding=utf-8
import tensorflow as tf
import numpy as np
import os 
from PIL import Image
 
def test():

    # 创建tfrecord文件
    train_filename = "train.tfrecords"
    if os.path.exists(train_filename):
        os.remove(train_filename)
 
   # 创建.tfrecord文件，准备写入
    writer = tf.python_io.TFRecordWriter('./'+train_filename) 


   #创建一个循环来依次写入数据
    for i in range(100):
        img_raw = np.random.randint(0,255,size=(300,300,3))       
        img_raw = img_raw.tostring()       # 转换为字符串   
        example = tf.train.Example(features=tf.train.Features(
                feature={
                'label': tf.train.Feature(int64_list = tf.train.Int64List(value=[i])),     
                'img_raw':tf.train.Feature(bytes_list = tf.train.BytesList(value=[img_raw]))
                }))
        writer.write(example.SerializeToString())  # 序列化保存
    writer.close()
    print ("保存tfrecord文件成功。")

    tfrecords_filename = "train.tfrecords"    
    filename_queue = tf.train.string_input_producer([tfrecords_filename],) #读入流中
    reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)   #返回文件名和文件
    features = tf.parse_single_example(serialized_example,
                                       features={
                                           'label': tf.FixedLenFeature([], tf.int64),
                                           'img_raw' : tf.FixedLenFeature([], tf.string),
                                       })  #取出包含image和label的feature对象
    image = tf.decode_raw(features['img_raw'],tf.int32)    
    image = tf.reshape(image, [300,300,3])
    label = tf.cast(features['label'], tf.int64)
    with tf.Session() as sess: #开始一个会话
        init_op = tf.global_variables_initializer()
        sess.run(init_op)

        coord=tf.train.Coordinator()  #创建一个协调器，管理线程
        threads= tf.train.start_queue_runners(coord=coord)  #启动QueueRunner, 此时文件名队列已经进队。
        for i in range(9):
            example, l = sess.run([image,label])             #在会话中取出image和label   
            print (example)         
            img=Image.fromarray(example, 'RGB')              #这里将array转为图像
            img.save('./'+str(i)+'_''Label_'+str(l)+'.jpg')  #存下图片           
            
        coord.request_stop()
        coord.join(threads)

def image_test():
	im = Image.open("E:/testdata/apple2orange/trainA/n07740461_2.jpg")  # 打开图像
	im.show()               # 显示图像
	img = np.array(im)      # image转 numpy
	img = img[:,:,0]        # 只取image的第2通道
	im=Image.fromarray(img) # numpy 转 image类
	im.show()               # 显示图像

if __name__ == '__main__':
	test()
	# image_test()

注意：

        coord=tf.train.Coordinator()  #创建一个协调器，管理线程
        threads= tf.train.start_queue_runners(coord=coord)  #启动QueueRunner, 此时文件名队列已经进队。

这两句非常重要，这两句实现的功能就是创建线程并使用QueueRunner对象来提取数据。简单来说：使用tf.train函数添加QueueRunner到tensorflow中。在运行任何训练步骤之前，需要调用tf.train.start_queue_runners函数，否则tensorflow将一直挂起。参考这里。tf.train.start_queue_runners 这个函数将会启动输入管道的线程，填充样本到队列中，以便出队操作可以从队列中拿到样本。这种情况下最好配合使用一个tf.train.Coordinator，这样可以在发生错误的情况下正确地关闭这些线程。如果你对训练迭代数做了限制，那么需要使用一个训练迭代数计数器，并且需要被初始化。

在tf.train中要创建这些队列和执行入队操作，就要添加QueueRunner到一个使用tf.train.add_queue_runner函数的数据流图中。每个QueueRunner负责一个阶段，处理那些需要在线程中运行的入队操作的列表。一旦数据流图构造成功，tf.train.start_queue_runners函数就会要求数据流图中每个QueueRunner去开始它的线程运行入队操作。
在执行训练的时候，队列会被后台的线程填充好。如果设置了最大训练迭代数（epoch），在某些时候，样本出队的操作可能会抛出一个tf.OutOfRangeError的错误。这是因为tensorflow的队列已经到达了最大实际的最大迭代数，没有更多可用的样本了。这也是为何推荐代码模板需要用try..except ..finally结构来处理这种错误。

扫描二维码关注公众号，回复： 2839519 查看本文章

上面代码读取的是单个的image和label，而在tensorflow训练时，一般是采取batch的方式去读入数据。tensorflow提供了两种方式，一种是shuffle_batch（tf.train.shuffle_batch），这种主要是用在训练中，随机选取样本组成batch。另外一种就是按照数据在tfrecord中的先后顺序生成batch（tf.train.batch，在Pix2Pix中采用这种方式）。这里采用tf.train.shuffle_batch方式：

image_batch, label_batch = tf.train.shuffle_batch(
      [image, label],
      batch_size=3,
      num_threads=4,
      capacity=10,
      min_after_dequeue=10)

参考：https://blog.csdn.net/happyhorizion/article/details/77894055

https://blog.csdn.net/qq_39037910/article/details/72900308

TensorFlow中的TFrecord

猜你喜欢