TensorFlow全新的數據讀取方式：Dataset API入門教程

知識 11-13

本文作者何之源，原文載於知乎專欄 AI Insight，AI研習社獲其授權發布。Dataset API是TensorFlow 1.3版本中引入的一個新的模塊，主要服務於數據讀取，構建輸入數據的pipeline。此前，在TensorFlow中讀取數據一般有兩種方法：使用placeholder讀內存中的數據使用queue讀硬碟中的數據（關於這種方式，可以參考我之前的一篇文章：十圖詳解tensorflow數據讀取機制（http://t.cn/RSSY56v））相Dataset API同時支持從內存和硬碟的讀取，相比之前的兩種方法在語法上更加簡潔易懂。此外，如果想要用到TensorFlow新出的Eager模式，就必須要使用Dataset API來讀取數據。本文就來為大家詳細地介紹一下Dataset API的使用方法（包括在非Eager模式和Eager模式下兩種情況）。Dataset API的導入在TensorFlow 1.3中，Dataset API是放在contrib包中的：

而在TensorFlow 1.4中，Dataset API已經從contrib包中移除，變成了核心API的一員：

下面的示例代碼將以TensorFlow 1.4版本為例，如果使用TensorFlow 1.3的話，需要進行簡單的修改（即加上contrib）。

基本概念：Dataset與Iterator

讓我們從基礎的類來了解Dataset API。參考Google官方給出的Dataset API中的類圖：

在初學時，我們只需要關注兩個最重要的基礎類：Dataset和Iterator。

Dataset可以看作是相同類型「元素」的有序列表。在實際使用時，單個「元素」可以是向量，也可以是字元串、圖片，甚至是tuple或者dict。

先以最簡單的，Dataset的每一個元素是一個數字為例：

這樣，我們就創建了一個dataset，這個dataset中含有5個元素，分別是1.0, 2.0, 3.0, 4.0, 5.0。

如何將這個dataset中的元素取出呢？方法是從Dataset中示例化一個Iterator，然後對Iterator進行迭代。

在非Eager模式下，讀取上述dataset中元素的方法為：

對應的輸出結果應該就是從1.0到5.0。語句iterator = dataset.make_one_shot_iterator()從dataset中實例化了一個Iterator，這個Iterator是一個「one shot iterator」，即只能從頭到尾讀取一次。one_element = iterator.get_next()表示從iterator里取出一個元素。由於這是非Eager模式，所以one_element只是一個Tensor，並不是一個實際的值。調用sess.run(one_element)後，才能真正地取出一個值。

如果一個dataset中元素被讀取完了，再嘗試sess.run(one_element)的話，就會拋出tf.errors.OutOfRangeError異常，這個行為與使用隊列方式讀取數據的行為是一致的。在實際程序中，可以在外界捕捉這個異常以判斷數據是否讀取完，請參考下面的代碼：

在Eager模式中，創建Iterator的方式有所不同。是通過tfe.Iterator(dataset)的形式直接創建Iterator并迭代。迭代時可以直接取出值，不需要使用sess.run()：

從內存中創建更複雜的Dataset之前我們用tf.data.Dataset.from_tensor_slices創建了一個最簡單的Dataset：

其實，tf.data.Dataset.from_tensor_slices的功能不止如此，它的真正作用是切分傳入Tensor的第一個維度，生成相應的dataset。

例如：

傳入的數值是一個矩陣，它的形狀為(5, 2)，tf.data.Dataset.from_tensor_slices就會切分它形狀上的第一個維度，最後生成的dataset中一個含有5個元素，每個元素的形狀是(2, )，即每個元素是矩陣的一行。

在實際使用中，我們可能還希望Dataset中的每個元素具有更複雜的形式，如每個元素是一個Python中的元組，或是Python中的詞典。例如，在圖像識別問題中，一個元素可以是{"image": image_tensor, "label": label_tensor}的形式，這樣處理起來更方便。

tf.data.Dataset.from_tensor_slices同樣支持創建這種dataset，例如我們可以讓每一個元素是一個詞典：

這時函數會分別切分"a"中的數值以及"b"中的數值，最終dataset中的一個元素就是類似於{"a": 1.0, "b": [0.9, 0.1]}的形式。

利用tf.data.Dataset.from_tensor_slices創建每個元素是一個tuple的dataset也是可以的：

對Dataset中的元素做變換：Transformation

Dataset支持一類特殊的操作：Transformation。一個Dataset通過Transformation變成一個新的Dataset。通常我們可以通過Transformation完成數據變換，打亂，組成batch，生成epoch等一系列操作。

常用的Transformation有：

map

batch

shuffle

repeat

下面就分別進行介紹。

（1）map

map接收一個函數，Dataset中的每個元素都會被當作這個函數的輸入，並將函數返回值作為新的Dataset，如我們可以對dataset中每個元素的值加1：

（2）batch

batch就是將多個元素組合成batch，如下面的程序將dataset中的每個元素組成了大小為32的batch：

（3）shuffle

shuffle的功能為打亂dataset中的元素，它有一個參數buffersize，表示打亂時使用的buffer的大小：

（4）repeat

repeat的功能就是將整個序列重複多次，主要用來處理機器學習中的epoch，假設原先的數據是一個epoch，使用repeat(5)就可以將之變成5個epoch：

如果直接調用repeat()的話，生成的序列就會無限重複下去，沒有結束，因此也不會拋出tf.errors.OutOfRangeError異常：

例子：讀入磁碟圖片與對應label

講到這裡，我們可以來考慮一個簡單，但同時也非常常用的例子：讀入磁碟中的圖片和圖片相應的label，並將其打亂，組成batch_size=32的訓練樣本。在訓練時重複10個epoch。

對應的程序為（從官方示常式序修改而來）：

在這個過程中，dataset經歷三次轉變：

運行dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))後，dataset的一個元素是(filename, label)。filename是圖片的文件名，label是圖片對應的標籤。

之後通過map，將filename對應的圖片讀入，並縮放為28x28的大小。此時dataset中的一個元素是(image_resized, label)

最後，dataset.shuffle(buffersize=1000).batch(32).repeat(10)的功能是：在每個epoch內將圖片打亂組成大小為32的batch，並重複10次。最終，dataset中的一個元素是(image_resized_batch, label_batch)，image_resized_batch的形狀為(32, 28, 28, 3)，而label_batch的形狀為(32, )，接下來我們就可以用這兩個Tensor來建立模型了。

Dataset的其它創建方法....

除了tf.data.Dataset.from_tensor_slices外，目前Dataset API還提供了另外三種創建Dataset的方式：

tf.data.TextLineDataset()：這個函數的輸入是一個文件的列表，輸出是一個dataset。dataset中的每一個元素就對應了文件中的一行。可以使用這個函數來讀入CSV文件。

tf.data.FixedLengthRecordDataset()：這個函數的輸入是一個文件的列表和一個record_bytes，之後dataset的每一個元素就是文件中固定位元組數record_bytes的內容。通常用來讀取以二進位形式保存的文件，如CIFAR10數據集就是這種形式。

tf.data.TFRecordDataset()：顧名思義，這個函數是用來讀TFRecord文件的，dataset中的每一個元素就是一個TFExample。

它們的詳細使用方法可以參閱文檔：Module: tf.data（http://t.cn/RlBVduF）

更多類型的Iterator....

在非Eager模式下，最簡單的創建Iterator的方法就是通過dataset.make_one_shot_iterator()來創建一個one shot iterator。除了這種one shot iterator外，還有三個更複雜的Iterator，即：

initializable iterator

reinitializable iterator

feedable iterator

initializable iterator必須要在使用前通過sess.run()來初始化。使用initializable iterator，可以將placeholder代入Iterator中，這可以方便我們通過參數快速定義新的Iterator。一個簡單的initializable iterator使用示例：

此時的limit相當於一個「參數」，它規定了Dataset中數的「上限」。

initializable iterator還有一個功能：讀入較大的數組。

在使用tf.data.Dataset.from_tensor_slices(array)時，實際上發生的事情是將array作為一個tf.constants保存到了計算圖中。當array很大時，會導致計算圖變得很大，給傳輸、保存帶來不便。這時，我們可以用一個placeholder取代這裡的array，並使用initializable iterator，只在需要時將array傳進去，這樣就可以避免把大數組保存在圖裡，示例代碼為（來自官方常式）：

reinitializable iterator和feedable iterator相比initializable iterator更複雜，也更加少用，如果想要了解它們的功能，可以參閱官方介紹（http://t.cn/RlBIUAu），這裡就不再贅述了。

總結

本文主要介紹了Dataset API的基本架構：Dataset類和Iterator類，以及它們的基礎使用方法。

在非Eager模式下，Dataset中讀出的一個元素一般對應一個batch的Tensor，我們可以使用這個Tensor在計算圖中構建模型。

在Eager模式下，Dataset建立Iterator的方式有所不同，此時通過讀出的數據就是含有值的Tensor，方便調試。

作為兼容兩種模式的Dataset API，在今後應該會成為TensorFlow讀取數據的主流方式。關於Dataset API的進一步介紹，可以參閱下面的資料：

Importing Data（http://t.cn/RlBMPHI）：官方Guide

Module: tf.data（http://t.cn/RlBVduF）： API文檔

Introduction to TensorFlow Datasets and Estimators（http://t.cn/RlBMpN0）：如何聯合使用Dataset和Estimator

新人福利

關注 AI 研習社（okweiwu），回復1領取

【超過 1000G 神經網路 / AI / 大數據，教程，論文】