轉載自https://zhuanlan.zhihu.com/p/32078191
出於想把程序放在美團雲的dls上跑的目的，學習了一下如何使用TF自帶的數據格式TFRecord（吐槽一下美團雲一次只能上傳20個文件的設定，想把一個數據集上傳上去真的太麻煩了）。說是針對分割問題的使用方法是因爲在分割問題裏label也是圖片，所以就只有圖片，其實TFRecord的使用方法非常靈活，只要能組織起自己的數據就可以隨便用起來。
TFRecord的格式
TFRecord文件中的數據都是通過tf.train.Example Protocol Buffer的格式存儲的：
message Example{
Features features = 1;
};
message Features{
map<string, Feature> feature = 1;
};
message Feature{
oneof kind {
BytesList bytes_list = 1;
FloatList float_list = 2;
Int64List int64_list = 3;
}
};
這個格式也是需要在程序中使用的，僅做了解。
將已有的數據集轉換成TFRecord格式
- 定義一個將已有的數據轉換成Feature數據結構的函數（官方教程中的函數）
-
將數據集中的文件名寫入一個txt文件中（爲了在程序中方便讀取），可以這樣來組織：
文件名1.jpg 文件名1.png
文件名2.jpg 文件名2.png
...... - 創建一個句柄來讀這個txt文件
- 定義一個TFRecord的writer
- 逐個文件來寫入TFRecoder文件
import tensorflow as tf
TXT_PATH = './dataset.txt'
TFRECORD_PATH = './dataset.tfrecord'
# 1.定義一個將已有的數據轉換成Feature數據結構的函數（官方教程中的函數）
def _bytes_feature(value):
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
# 3.創建一個句柄來讀這個txt文件
f = open(TXT_PATH)
# 4.定義一個TFRecord的writer
writer = tf.python_io.TFRecordWriter(TFRECORD_PATH)
# 5.逐個文件來寫入TFRecoder文件
for i in f.readlines():
# 每一行的文件名是用空格隔開的，所以需要使用split方法把string映射成list
item = i.split()
# 這裏可能有些教程裏會使用skimage或者opencv來讀取文件，但是我試了一下opencv的方法
# 一個400MB左右的訓練集先用cv2.imread轉換成numpy數組再轉成string，最後寫入TFRecord
# 得到的文件有17GB，所以這裏推薦使用FastGFile方法，生成的tfrecord文件會變小，
# 唯一不好的一點就是這種方法需要在讀取之後再進行一次解碼。
img = tf.gfile.FastGFile(item[0], 'rb').read()
label = tf.gfile.FastGFile(item[1], 'rb').read()
# 按照第一部分中Example Protocol Buffer的格式來定義要存儲的數據格式
example = tf.train.Example(features=tf.train.Features(feature={
'raw_image': _bytes_feature(img),
'label': _bytes_feature(label)
}))
# 最後將example寫入TFRecord文件
writer.write(example.SerializeToString())
writer.close()
f.close()#自己修改代碼# 1.定義一個將已有的數據轉換成Feature數據結構的函數（官方教程中的函數）
def _bytes_feature(value):
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
# 獲取文件列表
files_img = tf.gfile.Glob(os.path.join(images_dir,'*.jpg')) # 如我想列出images_dir下所有jpg文件路徑
files_anno = tf.gfile.Glob(os.path.join(annotations_dir,'*.png')) # 如我想列出annotations_dir下所有png文件路徑
# 4.定義一個TFRecord的writer
writer = tf.python_io.TFRecordWriter(TFRECORD_PATH)
# 5.逐個文件來寫入TFRecoder文件
for i in range(len(files_img)):
# 這裏可能有些教程裏會使用skimage或者opencv來讀取文件，但是我試了一下opencv的方法
# 一個400MB左右的訓練集先用cv2.imread轉換成numpy數組再轉成string，最後寫入TFRecord
# 得到的文件有17GB，所以這裏推薦使用FastGFile方法，生成的tfrecord文件會變小，
# 唯一不好的一點就是這種方法需要在讀取之後再進行一次解碼。
img = tf.gfile.FastGFile(files_img[i], 'rb').read()
label = tf.gfile.FastGFile(files_anno[i], 'rb').read()
# 按照第一部分中Example Protocol Buffer的格式來定義要存儲的數據格式
example = tf.train.Example(features=tf.train.Features(feature={
'raw_image': _bytes_feature(img),
'label': _bytes_feature(label)
}))
# 最後將example寫入TFRecord文件
writer.write(example.SerializeToString())
writer.close()因爲FastGFile讀取的是圖片沒有解碼過的的原始數據，所以在使用存在tfrecord中的這些原始數據時，需要對讀取出來的圖片原始數據進行解碼。
讀取TFRecord格式的數據
- 把所有的TFRecord文件名列表寫入隊列中（只有一個就寫一個文件名在列表中，多個就寫多個）
- 創建一個讀取器
- 將隊列中的tfrecord文件讀取爲example格式
- 根據定義數據的方式對應說明讀取的方式
- 對圖片進行解碼
# 1. 把所有的TFRecord文件名列表寫入隊列中（只有一個就寫一個文件名在列表中，多個就寫多個）
queue = tf.train.string_input_producer([TFRECORD_PATH], shuffle=True)
# 2. 創建一個讀取器
reader = tf.TFRecordReader()
# 3. 將隊列中的tfrecord文件讀取爲example格式
_, serialized_example = reader.read(queue)
# 4. 根據定義數據的方式對應說明讀取的方式
features = tf.parse_single_example(serialized_example,
features={
'raw_image': tf.FixedLenFeature([], tf.string)
'label': tf.FixedLenFeature([], tf.string)
})
img = features['raw_image']
label = features['label']
# 5. 對圖片進行解碼
img = tf.image_decode_jpeg(img, channels=3)
label = tf.image_decode_png(label, channels=1)
# 然後就可以用tf.train.batch方法來生成一個batch的image和label啦
img_batch, label_batch = tf.train.batch([img, label], BATCH_SIZE)
關於速度
因爲在沒有使用TFRecord時數據集是放在SSD上並且也是有解碼的步驟的，所以在使用了TFRecord之後數據的讀取沒有明顯變快，只是可以將上千張圖片的數據集壓縮到一個文件中這一點是比較方便的（拷一個巨大的數據集真的太慢了）。