Tfrecord檔案是tensorflow專門設計的一種訓練樣本儲存格式，將訓練樣本打包成tfrecord格式後能夠加快檔案的讀取效率。所以訓練網路的第一步就是將自己的訓練集樣本打包生成tfrecord格式。本文主要介紹兩種tfrecord打包方式，這兩種方式的主要區別在於生成的tfrecord檔案大小不同。

方式一：利用常用影象處理庫讀取影象並解碼，轉換成二進位制檔案進行儲存，網路上找到的基本上都是這種方式。

寫入tfrecord檔案

def data_to_tfrecord(images, labels, filename):  # images中儲存的是所有影象路徑的一個列表
    """ Save data into TFRecord """              # labels是images中每個影象對應的標籤
    if os.path.isfile(filename):                 # filename是tfrecord檔名稱
        print("%s exists" % filename)
        return
    print("Converting data into %s ..." % filename)
    writer = tf.python_io.TFRecordWriter(filename)
    for index, img_file in zip(labels, images):
        img1 = Image.open(img_file)   # 通過PIL包中的Images函式讀取、解碼圖片
        width, height = img1.size   # 獲取影象的寬、高參數
        img_raw = img1.tobytes()   # 將影象轉換成二進位制序列
        label = int(index)   # 圖片對應的標籤

        example = tf.train.Example(
            features=tf.train.Features(
                feature={
                    'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[label])),    # 儲存標籤
                    'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw])),  # 儲存二進位制序列
                    'img_width': tf.train.Feature(int64_list=tf.train.Int64List(value=[width])),  # 儲存影象的寬度
                    'img_height': tf.train.Feature(int64_list=tf.train.Int64List(value=[height]))  # 儲存影象的高
                }
            )
        )
        writer.write(example.SerializeToString())  # Serialize To String
    writer.close()
 

讀取tfrecord檔案

import numpy as np
import tensorflow as tf
import tensorlayer as tl


def read_and_decode(filename):
    """ Return tensor to read from TFRecord """
    filename_queue = tf.train.string_input_producer([filename])
    reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)
    features = tf.parse_single_example(
        serialized_example, features={
            'label': tf.FixedLenFeature([], tf.int64),   # 從tfrecord檔案中讀取各種資訊
            'img_raw': tf.FixedLenFeature([], tf.string),
            'img_width': tf.FixedLenFeature([], tf.int64),
            'img_height': tf.FixedLenFeature([], tf.int64)
        }
    )
    # You can do more image distortion here for training data
    width = tf.cast(features['img_width'], tf.int32)    # 轉型
    height = tf.cast(features['img_height'], tf.int32)
    img = tf.decode_raw(features['img_raw'], tf.uint8)   # 從二進位制檔案轉成uint8
    img = tf.reshape(img, [height, width, 3])   # 對影象進行reshape，注意在tfrecord檔案中儲存的是一序列，並沒有形狀
    img = tf.image.resize_images(img, [32, 32])   # 將影象統一到同一尺寸
    # img = tf.cast(img, tf.float32) #* (1. / 255) - 0.5
    label = tf.cast(features['label'], tf.int32)
    return img, label


# Example to visualize data
img, label = read_and_decode("train.tfrecord")
img_batch, label_batch = tf.train.shuffle_batch([img, label],
                                                batch_size=4,
                                                capacity=5000,
                                                min_after_dequeue=100,
                                                num_threads=1)
print("img_batch   : %s" % img_batch._shape)
print("label_batch : %s" % label_batch._shape)

init = tf.initialize_all_variables()
with tf.Session() as sess:
    sess.run(init)
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)

    for i in range(3):  # number of mini-batch (step)
        print("Step %d" % i)
        val, l = sess.run([img_batch, label_batch])
        # exit()
        print(val.shape, l)
        tl.visualize.images2d(val, second=1, saveable=False, name='batch'+str(i), dtype=np.uint8, fig_idx=2020121)

    coord.request_stop()
    coord.join(threads)
    sess.close()
 

方式二：利用tf.gfile.FastGFile讀取影象資訊（貌似並沒有解碼），轉換成二進位制檔案儲存。

這個方法是我在看tensorflow在github的slim框架中的生成tfrecord檔案所使用的方法。

寫入tfrecord檔案

def data_to_tfrecord(images, labels, filename):
    """ Save data into TFRecord """
    if os.path.isfile(filename):
        print("%s exists" % filename)
        return
    print("Converting data into %s ..." % filename)
    writer = tf.python_io.TFRecordWriter(filename)
    for index, img_file in zip(labels, images):
        img1 = Image.open(img_file)
        width, height = img1.size
        # img_raw = img1.tobytes()
        img_raw = tf.gfile.FastGFile(img_file, 'rb').read()  # 與方式一不同的是使用的FastGFile函式
        label = int(index)

        example = tf.train.Example(
            features=tf.train.Features(
                feature={
                    'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[label])),
                    'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw])),
                    'img_width': tf.train.Feature(int64_list=tf.train.Int64List(value=[width])),
                    'img_height': tf.train.Feature(int64_list=tf.train.Int64List(value=[height]))
                }
            )
        )
        writer.write(example.SerializeToString())  # Serialize To String
    writer.close()
 

讀取tfrecord檔案

import numpy as np
import tensorflow as tf
import tensorlayer as tl


def read_and_decode(filename):
    """ Return tensor to read from TFRecord """
    filename_queue = tf.train.string_input_producer([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),
            'img_width': tf.FixedLenFeature([], tf.int64),
            'img_height': tf.FixedLenFeature([], tf.int64)
        }
    )
    # You can do more image distortion here for training data
    width = tf.cast(features['img_width'], tf.int32)
    height = tf.cast(features['img_height'], tf.int32)
    # img = tf.decode_raw(features['img_raw'], tf.uint8)
    img = tf.image.decode_jpeg(features['img_raw'])  # 與方式一的不同點在於需要用decode_jpeg解碼
    img = tf.reshape(img, [height, width, 3])
    img = tf.image.resize_images(img, [32, 32])
    # img = tf.cast(img, tf.float32) #* (1. / 255) - 0.5
    label = tf.cast(features['label'], tf.int32)
    return img, label


# Example to visualize data
img, label = read_and_decode("train")
img_batch, label_batch = tf.train.shuffle_batch([img, label],
                                                batch_size=4,
                                                capacity=5000,
                                                min_after_dequeue=100,
                                                num_threads=1)
print("img_batch   : %s" % img_batch._shape)
print("label_batch : %s" % label_batch._shape)

init = tf.initialize_all_variables()
with tf.Session() as sess:
    sess.run(init)
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)

    for i in range(3):  # number of mini-batch (step)
        print("Step %d" % i)
        val, l = sess.run([img_batch, label_batch])
        # exit()
        print(val.shape, l)
        tl.visualize.images2d(val, second=1, saveable=False, name='batch'+str(i), dtype=np.uint8, fig_idx=2020121)

    coord.request_stop()
    coord.join(threads)
    sess.close()

兩種方式的區別

兩種方式雖然在程式碼上就那麼一兩行的區別，當對於生成的tfrecord檔案還是有很大的區別的。我用的同樣的影象樣本集，約200M左右，用方式一生成的tfrecord檔案約900M，用方式二生成的tfrecord檔案約200M。很明顯在佔用記憶體方面有著很大的區別。據我個人猜測，方案一將影象解碼後在轉成二進位制檔案的，方案二並沒有解碼而是直接轉成二進位制檔案進行儲存，所以在讀取時需要進行影象解碼。這僅是我個人猜測，如果有懂的大神，還望賜教。