文章目錄
爲什麼要使用 TFRecord 文件
正常情況下我們用於訓練的文件夾內部往往會存着成千上萬的圖片或文本等文件，這些文件通常被散列存放。這種存儲方式有一些缺點：
- 佔用磁盤空間；
- 在一個個讀取的時候會非常耗時；
- 佔用大量內存空間（有的大型數據不足以一次性加載）。
此時 TFRecord 格式的文件存儲形式會很合理的幫我們存儲數據。TFRecord 內部使用了 “Protocol Buffer” 二進制數據編碼方案，它只佔用一個內存塊，只需要一次性加載一個二進制文件的方式即可，簡單，快速，尤其對大型訓練數據很友好。而且當我們的訓練數據量比較大的時候，可以將數據分成多個 TFRecord 文件，來提高處理效率。
什麼是 TFRecord 文件
TFRecord 是 TensorFlow 中的數據集存儲格式。當我們將數據集整理成 TFRecord 格式後，TensorFlow 就可以高效地讀取和處理這些數據集，從而幫助我們更高效地進行大規模的模型訓練。
TFRecord 可以理解爲一系列序列化的 tf.train.Example 元素所組成的列表文件，而每一個 tf.train.Example 又由若干個 tf.train.Feature 的字典組成。形式如下：
[
{ # example 1 (tf.train.Example)
'feature_1': tf.train.Feature,
...
'feature_k': tf.train.Feature
},
...
{ # example N (tf.train.Example)
'feature_1': tf.train.Feature,
...
'feature_k': tf.train.Feature
}
]
怎樣寫入 TFRecord 文件
步驟：
- 讀取該數據元素到內存；
- 建立 tf.train.Feature 的字典；
- 將該元素轉換爲 tf.train.Example 對象（每一個 tf.train.Example 由若干個 tf.train.Feature 的字典組成）；
- 將該 tf.train.Example 對象序列化爲字符串，並通過一個預先定義的 tf.io.TFRecordWriter 寫入 TFRecord 文件。
1、導入需要的庫
import tensorflow as tf
import os
2、導入圖片
data_dir = './faces/' # 圖片所在文件夾
tfrecord_file = './tfrecord_files/train.tfrecords' # 要保存的 TFRecord 文件
train_filenames = [data_dir + filename for filename in os.listdir(data_dir)] # 將所有圖片的名稱寫入一個列表中
3、寫入 TFRecord 文件
with tf.io.TFRecordWriter(tfrecord_file) as writer:
for filename in train_filenames:
image = open(filename, 'rb').read() # 讀取數據集圖片到內存，image 爲一個 Byte 類型的字符串
feature = { # 建立 tf.train.Feature 字典
'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image])) # 圖片是一個 Bytes 對象
}
# tf.train.Example 在 tf.train.Features 外面又多了一層封裝
example = tf.train.Example(features=tf.train.Features(feature=feature)) # 通過字典建立 Example
writer.write(example.SerializeToString()) # 將 Example 序列化並寫入 TFRecord 文件
值得注意的是， tf.train.Feature 支持三種數據格式：
-
tf.train.BytesList ：字符串或原始 Byte 文件（如圖片），通過 bytes_list 參數傳入一個由字符串數組初始化的 tf.train.BytesList 對象；
-
tf.train.FloatList ：浮點數，通過 float_list 參數傳入一個由浮點數數組初始化的 tf.train.FloatList 對象；
-
tf.train.Int64List ：整數，通過 int64_list 參數傳入一個由整數數組初始化的 tf.train.Int64List 對象。
對這三種數據格式分別舉例說明：
# 將字符串列表轉化爲utf-8編碼
favorite_books = [name.encode('utf-8')
for name in ['machine learning', 'cc150']]
# 生成bytes_list
favorite_books_bytelist = tf.train.BytesList(value = favorite_books)
print(favorite_books_bytelist)
# 生成float_list
hours_floatlist = tf.train.FloatList(value = [15.5, 9.5, 7.0, 8.0])
print(hours_floatlist)
# 生成int64_list
age_int64list = tf.train.Int64List(value=[42]) # 如果只希望保存一個元素而非數組，傳入一個只有一個元素的數組即可
print(age_int64list)
value: "machine learning"
value: "cc150"
value: 15.5
value: 9.5
value: 7.0
value: 8.0
value: 42
運行以上代碼即可在 tfrecord_file 所指向的文件地址獲得一個 train.tfrecords 文件。
怎樣讀取 TFRecord 文件
-
通過 tf.data.TFRecordDataset 讀入原始的 TFRecord 文件（此時文件中的 tf.train.Example 對象尚未被反序列化），獲得一個 tf.data.Dataset 數據集對象；
-
通過 Dataset.map 方法，對該數據集對象中的每一個序列化的 tf.train.Example 字符串執行 tf.io.parse_single_example 函數，從而實現反序列化。
我們可以通過以下代碼，讀取之前建立的 train.tfrecords 文件。
1、初步讀取 TFRecord 文件
raw_dataset = tf.data.TFRecordDataset(tfrecord_file)
這裏得到的 raw_dataset 仍是字符串類型。
2、生成描述文件
feature_description = {
'image': tf.io.FixedLenFeature([], tf.string)
}
這裏的 feature_description 類似於一個數據集的 “描述文件”，通過一個由鍵值對組成的字典，告知解碼器每個 Feature 的類型是什麼。
tf.io.FixedLenFeature 的三個輸入參數 shape 、 dtype 和 default_value （可省略）爲每個 Feature 的形狀、類型和默認值。這裏我們的數據項都是單個的字符串，所以 shape 爲空數組。
3、定義解碼器
def _parse_example(example_string): # 將 TFRecord 文件中的每一個序列化的 tf.train.Example 解碼
feature_dict = tf.io.parse_single_example(example_string, feature_description)
feature_dict['image'] = tf.io.decode_jpeg(feature_dict['image']) # 解碼 JEPG 圖片
return feature_dict['image']
dataset = raw_dataset.map(_parse_example)
得到的數據集對象 dataset 是一個可以用於訓練的 tf.data.Dataset 對象。
4、展示圖片
import matplotlib.pyplot as plt
%matplotlib inline
for image in dataset.take(1):
plt.title('face')
plt.imshow(image.numpy())
plt.show()
