以Cifar10图像数据为例详解tensorflow对tfrecord数据的读写操作

末蓝、 2023-06-24 13:28 2阅读 0赞

> **陈伟@航天科技智慧城市研究院**chenwei@ascs.tech

## tfrecord格式说明 ##

正常情况下我们训练文件夹经常会生成 train, test 或者val文件夹，这些文件夹内部往往会存着成千上万的图片或文本等文件，这些文件被散列存着，这样不仅占用磁盘空间，并且再被一个个读取的时候会非常慢，繁琐。占用大量内存空间（有的大型数据不足以一次性加载）。此时我们TFRecord格式的文件存储形式会很合理的帮我们存储数据。TFRecord内部使用了“Protocol Buffer”二进制数据编码方案，它只占用一个内存块，只需要一次性加载一个二进制文件的方式即可，简单，快速，尤其对大型训练数据很友好。而且当我们的训练数据量比较大的时候，可以将数据分成多个TFRecord文件，来提高处理效率。

## 前置工作 ##

我们在[《cifar10数据通过tensorflow转换成图像》][cifar10_tensorflow]将cifar10数据转换成了图像数据。本文将演示将这些图片序列化成tfrecord数据，以便后续的训练。

## 写tfrecord数据writer\_cifar10.py ##

import tensorflow as tf
    import cv2
    import numpy as np
    import glob
    
    classification = ['airplane',
                      'automobile',
                      'bird',
                      'cat',
                      'deer',
                      'dog',
                      'frog',
                      'horse',
                      'ship',
                      'truck']
    
    
    idx = 0
    im_data = []  # 所有类别的图片文件
    im_labels = []  # 所有文件的标签
    for path in classification:  # 对于每个类别
        path = "data/image/train/" + path  # 形如：data/image/train/bird
        im_list = glob.glob(path + "/*")  # 获取这个类别下所有文件
        # Python在方括号中使用for循环，类似[0 for i in range(10)]，叫 列表解析List Comprehensions
        # 根据已有列表，高效创建新列表的方式。
        # 列表解析是Python迭代机制的一种应用，它常用于实现创建新的列表，因此用在[]中。
        # [expression for iter_val in iterable]
        # [expression for iter_val in iterable if cond_expr]
        im_label = [idx for i in range(im_list.__len__())]  # idx就只有0~9
        idx += 1  # 每次换类别加1
        im_data += im_list  # 把本次类别下的图片列表加到全部的图片列表中
        im_labels += im_label  # 把本次的标签加到全部的标签列表中
    
    print(im_labels)
    # print(im_data)
    
    # 下面生成.tfrecord文件
    tfrecord_file = "data/train.tfrecord"
    # 利用TFRecordWriter写入到文件中
    writer = tf.compat.v1.python_io.TFRecordWriter(tfrecord_file)
    
    index = [i for i in range(im_data.__len__())]
    
    # 打乱图片的顺序
    np.random.shuffle(index)
    
    for i in range(im_data.__len__()):
        im_d = im_data[index[i]]
        im_l = im_labels[index[i]]
        # opencv对图片进行读取
        data = cv2.imread(im_d)
        # 也可以通过这种方式读取图片数据，本身就是byte格式的
        # data = tf.gfile.FastGFile(im_d, "rb").read()
        ex = tf.train.Example(
            # Features是用于描述机器学习模型训练或推理的特征的协议消息，用键值对表示数据。
            # 一个Features中包括可能包含零个或多个值的列表。 这些列表是基本值BytesList，FloatList，Int64List。Feature按名称分类。 Feature的消息包含从名称到功能的映射。
            features=tf.train.Features(
                feature={ 
                    "image": tf.train.Feature(  # 图像数据使用byte类型进行存储
                        bytes_list=tf.train.BytesList(
                            value=[data.tobytes()])),
                    "label": tf.train.Feature(  # 标签数据是int类型的数据
                        int64_list=tf.train.Int64List(
                            value=[im_l])),
                }
            )
        )
        # 封装好的数据进行序列化，并且写入tfrecord中
        writer.write(ex.SerializeToString())
    # 关闭writer
    writer.close()

代码已经上传到[Github][]

## 读tfrecord数据reader\_cifar10.py ##

import tensorflow as tf
    
    # 文件路径列表
    filelist = ['data/train.tfrecord']
    # 定义文件队列
    file_queue = tf.train.string_input_producer(filelist,
                                                num_epochs=None,
                                                shuffle=True)
    # 通过TFRcordReader对象读取tfrecord文件
    reader = tf.TFRecordReader()
    # 返回文件名和文件
    _, ex = reader.read(file_queue)
    
    # 需要对序列化的数据进行解码
    feature = { 
        'image': tf.FixedLenFeature([], tf.string),  # byte型的解码成string型
        'label': tf.FixedLenFeature([], tf.int64)
    }
    
    batchsize = 1
    
    batch = tf.train.shuffle_batch([ex], batchsize, capacity=batchsize*10,
                           min_after_dequeue=batchsize*5)
    
    # 反序列化数据
    example = tf.parse_example(batch, features=feature)
    
    image = example['image']
    label = example['label']
    
    # 对byte数据解码成uint8类型的数据
    image = tf.decode_raw(image, tf.uint8)
    # 需要reshape，否则是一个向量
    image = tf.reshape(image, [-1, 32, 32, 3])
    
    with tf.Session() as sess:
        # 线程的协调器
        coord = tf.train.Coordinator()
    
        sess.run(tf.local_variables_initializer())
        threads = tf.train.start_queue_runners(sess, coord)
    
        for i in range(1):
            # image_bth, _ = sess.run([image, label])
            # import cv2
            # cv2.imshow("image", image_bth[0, ...])
            # cv2.waitKey(0)
            print(sess.run(label))  # 会打印一个batchsize数量的lable
    
        # 请求线程结束
        coord.request_stop()
        # 等待线程终止
        coord.join(threads)

*  关于tf.train.slice\_input\_producer和tf.train.string\_input\_producer张量及文件生成器已经在文章[《详解tensorflow的tf.train.slice\_input\_producer和tf.train.string\_input\_producer生成器》][tensorflow_tf.train.slice_input_producer_tf.train.string_input_producer]详细说明
 *  tf.train.shuffle\_batch：在文章[《详解tensorflow之tf.train.batch与tf.train.shuffle\_batch附完整代码》][tensorflow_tf.train.batch_tf.train.shuffle_batch]有详细说明
 *  reader\_cifar10.py代码已上传到[Github][Github 1]

[cifar10_tensorflow]: https://blog.csdn.net/keyandi/article/details/103667443
[Github]: https://github.com/cwyd0822/cifar10-tensorflow-read-write/blob/master/writer_cifar10.py
[tensorflow_tf.train.slice_input_producer_tf.train.string_input_producer]: https://blog.csdn.net/keyandi/article/details/103683761
[tensorflow_tf.train.batch_tf.train.shuffle_batch]: https://blog.csdn.net/keyandi/article/details/103685935
[Github 1]: https://github.com/cwyd0822/cifar10-tensorflow-read-write/blob/master/reader_cifar10.py