目錄
數據集對象的建立
提示
數據集對象的預處理
`Dataset.shuffle()` 時緩沖區大小 `buffer_size` 的設置
使用 `tf.data` 的並行化策略提高訓練流程效率
數據集元素的獲取與使用
實例：cats_vs_dogs 圖像分類
很多時候，我們希望使用自己的數據集來訓練模型。然而，面對一堆格式不一的原始數據文件，將其預處理並讀入程序的過程往往十分繁瑣，甚至比模型的設計還要耗費精力。比如，為了讀入一批圖像文件，我們可能需要糾結於 python 的各種圖像處理包(比如 pillow )，自己設計 Batch 的生成方式，最后還可能在運行的效率上不盡如人意。為此，TensorFlow 提供了 tf.data 這一模塊，包括了一套靈活的數據集構建 API，能夠幫助我們快速、高效地構建數據輸入的流水線，尤其適用於數據量巨大的場景。
數據集對象的建立
tf.data 的核心是 tf.data.Dataset 類，提供了對數據集的高層封裝。tf.data.Dataset 由一系列的可迭代訪問的元素(element)組成，每個元素包含一個或多個張量。比如說，對於一個由圖像組成的數據集，每個元素可以是一個形狀為 長×寬×通道數 的圖片張量，也可以是由圖片張量和圖片標簽張量組成的元組(Tuple)。
最基礎的建立 tf.data.Dataset 的方法是使用 tf.data.Dataset.from_tensor_slices() ，適用於數據量較小(能夠整個裝進內存)的情況。具體而言，如果我們的數據集中的所有元素通過張量的第 0 維，拼接成一個大的張量(例如，前節的 MNIST 數據集的訓練集即為一個 [60000, 28, 28, 1] 的張量，表示了 60000 張 28*28 的單通道灰度圖像)，那么我們提供一個這樣的張量或者第 0 維大小相同的多個張量作為輸入，即可按張量的第 0 維展開來構建數據集，數據集的元素數量為張量第 0 維的大小。具體示例如下：
import tensorflow as tf
import numpy as np
X = tf.constant([2013, 2014, 2015, 2016, 2017])
Y = tf.constant([12000, 14000, 15000, 16500, 17500])
# 也可以使用NumPy數組，效果相同
# X = np.array([2013, 2014, 2015, 2016, 2017])
# Y = np.array([12000, 14000, 15000, 16500, 17500])
dataset = tf.data.Dataset.from_tensor_slices((X, Y))
for x, y in dataset:
print(x.numpy(), y.numpy())
輸出:
2013 12000
2014 14000
2015 15000
2016 16500
2017 17500
當提供多個張量作為輸入時，張量的第 0 維大小必須相同，且必須將多個張量作為元組(Tuple，即使用 Python 中的小括號)拼接並作為輸入。
類似地，我們可以載入前章的 MNIST 數據集：
import matplotlib.pyplot as plt
(train_data, train_label), (_, _) = tf.keras.datasets.mnist.load_data()
train_data = np.expand_dims(train_data.astype(np.float32) / 255.0, axis=-1) # [60000, 28, 28, 1]
mnist_dataset = tf.data.Dataset.from_tensor_slices((train_data, train_label))
for image, label in mnist_dataset:
plt.title(label.numpy())
plt.imshow(image.numpy()[:, :, 0])
plt.show()
輸出
提示
TensorFlow Datasets 提供了一個基於 tf.data.Datasets 的開箱即用的數據集集合，相關內容可參考 TensorFlow Datasets 。例如，使用以下語句：
import tensorflow_datasets as tfds
dataset = tfds.load("mnist", split=tfds.Split.TRAIN, as_supervised=True)
即可快速載入 MNIST 數據集。
對於特別巨大而無法完整載入內存的數據集，我們可以先將數據集處理為 TFRecord 格式，然后使用 tf.data.TFRocordDataset() 進行載入。詳情請參考后文
數據集對象的預處理
tf.data.Dataset 類為我們提供了多種數據集預處理方法。最常用的如：
Dataset.map(f) ：對數據集中的每個元素應用函數 f ，得到一個新的數據集(這部分往往結合 tf.io 進行讀寫和解碼文件， tf.image 進行圖像處理);
Dataset.shuffle(buffer_size) ：將數據集打亂(設定一個固定大小的緩沖區(Buffer)，取出前 buffer_size 個元素放入，並從緩沖區中隨機采樣，采樣后的數據用后續數據替換);
Dataset.batch(batch_size) ：將數據集分成批次，即對每 batch_size 個元素，使用 tf.stack() 在第 0 維合並，成為一個元素;
除此以外，還有 Dataset.repeat() (重復數據集的元素)、 Dataset.reduce() (與 Map 相對的聚合操作)、 Dataset.take() (截取數據集中的前若干個元素)等，可參考 API 文檔 進一步了解。
以下以 MNIST 數據集進行示例。
使用 Dataset.map() 將所有圖片旋轉 90 度：
def rot90(image, label):
image = tf.image.rot90(image)
return image, label
mnist_dataset = mnist_dataset.map(rot90)
for image, label in mnist_dataset:
plt.title(label.numpy())
plt.imshow(image.numpy()[:, :, 0])
plt.show()
輸出
使用 Dataset.batch() 將數據集划分批次，每個批次的大小為 4：
mnist_dataset = mnist_dataset.batch(4)
for images, labels in mnist_dataset: # image: [4, 28, 28, 1], labels: [4]
fig, axs = plt.subplots(1, 4)
for i in range(4):
axs[i].set_title(labels.numpy()[i])
axs[i].imshow(images.numpy()[i, :, :, 0])
plt.show()
輸出
使用 Dataset.shuffle() 將數據打散后再設置批次，緩存大小設置為 10000：
mnist_dataset = mnist_dataset.shuffle(buffer_size=10000).batch(4)
for images, labels in mnist_dataset:
fig, axs = plt.subplots(1, 4)
for i in range(4):
axs[i].set_title(labels.numpy()[i])
axs[i].imshow(images.numpy()[i, :, :, 0])
plt.show()
輸出
第一次運行
第二次運行
可見每次的數據都會被隨機打散。
Dataset.shuffle() 時緩沖區大小 buffer_size 的設置
tf.data.Dataset 作為一個針對大規模數據設計的迭代器，本身無法方便地獲得自身元素的數量或隨機訪問元素。因此，為了高效且較為充分地打散數據集，需要一些特定的方法。Dataset.shuffle() 采取了以下方法：
設定一個固定大小為 buffer_size 的緩沖區(Buffer);
初始化時，取出數據集中的前 buffer_size 個元素放入緩沖區;
每次需要從數據集中取元素時，即從緩沖區中隨機采樣一個元素並取出，然后從后續的元素中取出一個放回到之前被取出的位置，以維持緩沖區的大小。
因此，緩沖區的大小需要根據數據集的特性和數據排列順序特點來進行合理的設置。比如：
當 buffer_size 設置為 1 時，其實等價於沒有進行任何打散;
當數據集的標簽順序分布極為不均勻(例如二元分類時數據集前 N 個的標簽為 0，后 N 個的標簽為 1 時，較小的緩沖區大小會使得訓練時取出的 Batch 數據很可能全為同一標簽，從而影響訓練效果。一般而言，數據集的順序分布若較為隨機，則緩沖區的大小可較小，否則則需要設置較大的緩沖區。
使用 tf.data 的並行化策略提高訓練流程效率
當訓練模型時，我們希望充分利用計算資源，減少 CPU/GPU 的空載時間。然而有時，數據集的准備處理非常耗時，使得我們在每進行一次訓練前都需要花費大量的時間准備待訓練的數據，而此時 GPU 只能空載而等待數據，造成了計算資源的浪費，如下圖所示：
常規訓練流程，在准備數據時，GPU 只能空載。
此時， tf.data 的數據集對象為我們提供了 Dataset.prefetch() 方法，使得我們可以讓數據集對象 Dataset 在訓練時預取出若干個元素，使得在 GPU 訓練的同時 CPU 可以准備數據，從而提升訓練流程的效率，如下圖所示：
Dataset.prefetch() 的使用方法和前節的 Dataset.batch() 、 Dataset.shuffle() 等非常類似。繼續以前節的 MNIST 數據集為例，若希望開啟預加載數據，使用如下代碼即可：
mnist_dataset = mnist_dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
此處參數 buffer_size 既可手工設置，也可設置為 tf.data.experimental.AUTOTUNE 從而由 TensorFlow 自動選擇合適的數值。
與此類似， Dataset.map() 也可以利用多 GPU 資源，並行化地對數據項進行變換，從而提高效率。以前節的 MNIST 數據集為例，假設用於訓練的計算機具有 2 核的 CPU，我們希望充分利用多核心的優勢對數據進行並行化變換(比如前節的旋轉 90 度函數 rot90 )，可以使用以下代碼：
mnist_dataset = mnist_dataset.map(map_func=rot90, num_parallel_calls=2)
其運行過程如下圖所示：
通過設置 Dataset.map() 的 num_parallel_calls 參數實現數據轉換的並行化。上部分是未並行化的圖示，下部分是 2 核並行的圖示。
當然，這里同樣可以將 num_parallel_calls 設置為 tf.data.experimental.AUTOTUNE 以讓 TensorFlow 自動選擇合適的數值。
除此以外，還有很多提升數據集處理性能的方式，可參考 TensorFlow 文檔 進一步了解。
數據集元素的獲取與使用
構建好數據並預處理后，我們需要從其中迭代獲取數據以用於訓練。tf.data.Dataset 是一個 Python 的可迭代對象，因此可以使用 For 循環迭代獲取數據，即：
dataset = tf.data.Dataset.from_tensor_slices((A, B, C, ...))
for a, b, c, ... in dataset:
# 對張量a, b, c等進行操作，例如送入模型進行訓練
也可以使用 iter() 顯式創建一個 Python 迭代器並使用 next() 獲取下一個元素，即：
dataset = tf.data.Dataset.from_tensor_slices((A, B, C, ...))
it = iter(dataset)
a_0, b_0, c_0, ... = next(it)
a_1, b_1, c_1, ... = next(it)
Keras 支持使用 tf.data.Dataset 直接作為輸入。當調用 tf.keras.Model 的 fit() 和 evaluate() 方法時，可以將參數中的輸入數據 x 指定為一個元素格式為 (輸入數據, 標簽數據) 的 Dataset ，並忽略掉參數中的標簽數據 y 。例如，對於上述的 MNIST 數據集，常規的 Keras 訓練方式是：
model.fit(x=train_data, y=train_label, epochs=num_epochs, batch_size=batch_size)
使用 tf.data.Dataset 后，我們可以直接傳入 Dataset：
model.fit(mnist_dataset, epochs=num_epochs)
由於已經通過 Dataset.batch() 方法划分了數據集的批次，所以這里也無需提供批次的大小。
實例：cats_vs_dogs 圖像分類
以下代碼以貓狗圖片二分類任務為示例，展示了使用 tf.data 結合 tf.io 和 tf.image 建立 tf.data.Dataset 數據集，並進行訓練和測試的完整過程。使用前須將數據集解壓到代碼中 data_dir 所設置的目錄.
import tensorflow as tf
import os
num_epochs = 10
batch_size = 32
learning_rate = 0.001
data_dir = 'C:/datasets/cats_vs_dogs'
train_cats_dir = data_dir + '/train/cats/'
train_dogs_dir = data_dir + '/train/dogs/'
test_cats_dir = data_dir + '/valid/cats/'
test_dogs_dir = data_dir + '/valid/dogs/'
def _decode_and_resize(filename, label):
image_string = tf.io.read_file(filename) # 讀取原始文件
image_decoded = tf.image.decode_jpeg(image_string) # 解碼JPEG圖片
image_resized = tf.image.resize(image_decoded, [256, 256]) / 255.0
return image_resized, label
if __name__ == '__main__':
# 構建訓練數據集大連婦科醫院 http://xmobile.bhbyby.com/
train_cat_filenames = tf.constant([train_cats_dir + filename for filename in os.listdir(train_cats_dir)])
train_dog_filenames = tf.constant([train_dogs_dir + filename for filename in os.listdir(train_dogs_dir)])
train_filenames = tf.concat([train_cat_filenames, train_dog_filenames], axis=-1)
train_labels = tf.concat([
tf.zeros(train_cat_filenames.shape, dtype=tf.int32),
tf.ones(train_dog_filenames.shape, dtype=tf.int32)],
axis=-1)
train_dataset = tf.data.Dataset.from_tensor_slices((train_filenames, train_labels))
train_dataset = train_dataset.map(
map_func=_decode_and_resize,
num_parallel_calls=tf.data.experimental.AUTOTUNE)
# 取出前buffer_size個數據放入buffer，並從其中隨機采樣，采樣后的數據用后續數據替換
train_dataset = train_dataset.shuffle(buffer_size=23000)
train_dataset = train_dataset.batch(batch_size)
train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)
model = tf.keras.Sequential([
tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(256, 256, 3)),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Conv2D(32, 5, activation='relu'),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(2, activation='softmax')
])
model.compile(
optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
loss=tf.keras.losses.sparse_categorical_crossentropy,
metrics=[tf.keras.metrics.sparse_categorical_accuracy]
)
model.fit(train_dataset, epochs=num_epochs)
使用以下代碼進行測試：
# 構建測試數據集
test_cat_filenames = tf.constant([test_cats_dir + filename for filename in os.listdir(test_cats_dir)])
test_dog_filenames = tf.constant([test_dogs_dir + filename for filename in os.listdir(test_dogs_dir)])
test_filenames = tf.concat([test_cat_filenames, test_dog_filenames], axis=-1)
test_labels = tf.concat([
tf.zeros(test_cat_filenames.shape, dtype=tf.int32),
tf.ones(test_dog_filenames.shape, dtype=tf.int32)],
axis=-1)
test_dataset = tf.data.Dataset.from_tensor_slices((test_filenames, test_labels))
test_dataset = test_dataset.map(_decode_and_resize)
test_dataset = test_dataset.batch(batch_size)
print(model.metrics_names)
print(model.evaluate(test_dataset))
通過對以上示例進行性能測試，我們可以感受到 tf.data 的強大並行化性能。通過 prefetch() 的使用和在 map() 過程中加入 num_parallel_calls 參數，模型訓練的時間可縮減至原來的一半甚至更低。得出結果。