Dataset API是TensorFlow 1.3版本中引入的一個新的模塊，主要服務(wù)于數(shù)據(jù)讀取妻味，構(gòu)建輸入數(shù)據(jù)的pipeline。

此前，在TensorFlow中讀取數(shù)據(jù)一般有兩種方法：

• 使用placeholder讀內(nèi)存中的數(shù)據(jù)
• 使用queue讀硬盤中的數(shù)據(jù)（關(guān)于這種方式轧抗，可以參考我之前的一篇文章：十圖詳解tensorflow數(shù)據(jù)讀取機(jī)制）

相Dataset API同時(shí)支持從內(nèi)存和硬盤的讀取，相比之前的兩種方法在語法上****更加簡潔易懂瞬测。此外横媚，如果想要用到TensorFlow新出的Eager模式纠炮，就必須要使用Dataset API來讀取數(shù)據(jù)。

本文就來為大家詳細(xì)地介紹一下Dataset API的使用方法（包括在非Eager模式和Eager模式下兩種情況）灯蝴。

Dataset API的導(dǎo)入

在TensorFlow 1.3中恢口，Dataset API是放在contrib包中的：

tf.contrib.data.Dataset

而在TensorFlow 1.4中，Dataset API已經(jīng)從contrib包中移除穷躁，變成了核心API的一員：

tf.data.Dataset

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

基本概念：Dataset與Iterator

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

image

</figure>

在初學(xué)時(shí)，我們只需要關(guān)注兩個最重要的基礎(chǔ)類：Dataset和Iterator狡忙。

Dataset可以看作是相同類型“元素”的有序列表梳虽。在實(shí)際使用時(shí)，單個“元素”可以是向量灾茁，也可以是字符串窜觉、圖片，甚至是tuple或者dict删顶。

先以最簡單的竖螃，Dataset的每一個元素是一個數(shù)字為例：

import tensorflow as tf
import numpy as np

dataset = tf.data.Dataset.from_tensor_slices(np.array([1.0, 2.0, 3.0, 4.0, 5.0]))

這樣，我們就創(chuàng)建了一個dataset逗余，這個dataset中含有5個元素特咆，分別是1.0, 2.0, 3.0, 4.0, 5.0。

如何將這個dataset中的元素取出呢录粱？方法是從Dataset中示例化一個Iterator腻格，然后對Iterator進(jìn)行迭代。

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

iterator = dataset.make_one_shot_iterator()
one_element = iterator.get_next()
with tf.Session() as sess:
    for i in range(5):
        print(sess.run(one_element))

對應(yīng)的輸出結(jié)果應(yīng)該就是從1.0到5.0菜职。語句iterator = dataset.make_one_shot_iterator()從dataset中實(shí)例化了一個Iterator，這個Iterator是一個“one shot iterator”旗闽，即只能從頭到尾讀取一次酬核。one_element = iterator.get_next()表示從iterator里取出一個元素。由于這是非Eager模式适室，所以one_element只是一個Tensor嫡意，并不是一個實(shí)際的值。調(diào)用sess.run(one_element)后捣辆，才能真正地取出一個值蔬螟。

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

dataset = tf.data.Dataset.from_tensor_slices(np.array([1.0, 2.0, 3.0, 4.0, 5.0]))
iterator = dataset.make_one_shot_iterator()
one_element = iterator.get_next()
with tf.Session() as sess:
    try:
        while True:
            print(sess.run(one_element))
    except tf.errors.OutOfRangeError:
        print("end!")

在Eager模式中坎怪，創(chuàng)建Iterator的方式有所不同。是通過tfe.Iterator(dataset)的形式直接創(chuàng)建Iterator并迭代廓握。迭代時(shí)可以直接取出值芋忿，不需要使用sess.run()：

import tensorflow.contrib.eager as tfe
tfe.enable_eager_execution()

dataset = tf.data.Dataset.from_tensor_slices(np.array([1.0, 2.0, 3.0, 4.0, 5.0]))

for one_element in tfe.Iterator(dataset):
    print(one_element)

從內(nèi)存中創(chuàng)建更復(fù)雜的Dataset

之前我們用tf.data.Dataset.from_tensor_slices創(chuàng)建了一個最簡單的Dataset：

dataset = tf.data.Dataset.from_tensor_slices(np.array([1.0, 2.0, 3.0, 4.0, 5.0]))

其實(shí)，tf.data.Dataset.from_tensor_slices的功能不止如此疾棵，它的真正作用是切分傳入Tensor的第一個維度，生成相應(yīng)的dataset痹仙。

例如：

dataset = tf.data.Dataset.from_tensor_slices(np.random.uniform(size=(5, 2)))

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

在實(shí)際使用中恩溅，我們可能還希望Dataset中的每個元素具有更復(fù)雜的形式，如每個元素是一個Python中的元組谓娃，或是Python中的詞典脚乡。例如，在圖像識別問題中滨达，一個元素可以是{"image": image_tensor, "label": label_tensor}的形式奶稠，這樣處理起來更方便。

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

dataset = tf.data.Dataset.from_tensor_slices(
    {
        "a": np.array([1.0, 2.0, 3.0, 4.0, 5.0]),                                       
        "b": np.random.uniform(size=(5, 2))
    }
)

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

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

dataset = tf.data.Dataset.from_tensor_slices(
  (np.array([1.0, 2.0, 3.0, 4.0, 5.0]), np.random.uniform(size=(5, 2)))
)

對Dataset中的元素做變換：Transformation

Dataset支持一類特殊的操作：Transformation画株。一個Dataset通過Transformation變成一個新的Dataset辆飘。通常我們可以通過Transformation完成數(shù)據(jù)變換，打亂谓传，組成batch蜈项，生成epoch等一系列操作。

常用的Transformation有：

• map
• batch
• shuffle
• repeat

下面就分別進(jìn)行介紹良拼。

（1）map

map接收一個函數(shù)战得，Dataset中的每個元素都會被當(dāng)作這個函數(shù)的輸入，并將函數(shù)返回值作為新的Dataset庸推，如我們可以對dataset中每個元素的值加1：

dataset = tf.data.Dataset.from_tensor_slices(np.array([1.0, 2.0, 3.0, 4.0, 5.0]))
dataset = dataset.map(lambda x: x + 1) # 2.0, 3.0, 4.0, 5.0, 6.0

（2）batch

batch就是將多個元素組合成batch常侦，如下面的程序?qū)ataset中的每個元素組成了大小為32的batch：

dataset = dataset.batch(32)

（3）shuffle

shuffle的功能為打亂dataset中的元素浇冰，它有一個參數(shù)buffersize，表示打亂時(shí)使用的buffer的大辛觥：

dataset = dataset.shuffle(buffer_size=10000)

（4）repeat

repeat的功能就是將整個序列重復(fù)多次肘习，主要用來處理機(jī)器學(xué)習(xí)中的epoch，假設(shè)原先的數(shù)據(jù)是一個epoch坡倔，使用repeat(5)就可以將之變成5個epoch：

dataset = dataset.repeat(5)

如果直接調(diào)用repeat()的話漂佩，生成的序列就會無限重復(fù)下去，沒有結(jié)束罪塔，因此也不會拋出tf.errors.OutOfRangeError異常：

dataset = dataset.repeat()

例子：讀入磁盤圖片與對應(yīng)label

講到這里投蝉，我們可以來考慮一個簡單，但同時(shí)也非常常用的例子：讀入磁盤中的圖片和圖片相應(yīng)的label征堪，并將其打亂瘩缆，組成batch_size=32的訓(xùn)練樣本。在訓(xùn)練時(shí)重復(fù)10個epoch佃蚜。

對應(yīng)的程序?yàn)椋◤墓俜绞纠绦蛐薷亩鴣恚?/p>

# 函數(shù)的功能時(shí)將filename對應(yīng)的圖片文件讀進(jìn)來庸娱，并縮放到統(tǒng)一的大小
def _parse_function(filename, label):
  image_string = tf.read_file(filename)
  image_decoded = tf.image.decode_image(image_string)
  image_resized = tf.image.resize_images(image_decoded, [28, 28])
  return image_resized, label

# 圖片文件的列表
filenames = tf.constant(["/var/data/image1.jpg", "/var/data/image2.jpg", ...])
# label[i]就是圖片filenames[i]的label
labels = tf.constant([0, 37, ...])

# 此時(shí)dataset中的一個元素是(filename, label)
dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))

# 此時(shí)dataset中的一個元素是(image_resized, label)
dataset = dataset.map(_parse_function)

# 此時(shí)dataset中的一個元素是(image_resized_batch, label_batch)
dataset = dataset.shuffle(buffersize=1000).batch(32).repeat(10)

在這個過程中，dataset經(jīng)歷三次轉(zhuǎn)變：

• 運(yùn)行dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))后谐算，dataset的一個元素是(filename, label)熟尉。filename是圖片的文件名，label是圖片對應(yīng)的標(biāo)簽洲脂。
• 之后通過map斤儿，將filename對應(yīng)的圖片讀入，并縮放為28x28的大小腮考。此時(shí)dataset中的一個元素是(image_resized, label)
• 最后雇毫，dataset.shuffle(buffersize=1000).batch(32).repeat(10)的功能是：在每個epoch內(nèi)將圖片打亂組成大小為32的batch，并重復(fù)10次踩蔚。最終棚放，dataset中的一個元素是(image_resized_batch, label_batch)，image_resized_batch的形狀為(32, 28, 28, 3)馅闽，而label_batch的形狀為(32, )飘蚯，接下來我們就可以用這兩個Tensor來建立模型了。

Dataset的其它創(chuàng)建方法....

除了tf.data.Dataset.from_tensor_slices外福也，目前Dataset API還提供了另外三種創(chuàng)建Dataset的方式：

• tf.data.TextLineDataset()：這個函數(shù)的輸入是一個文件的列表局骤，輸出是一個dataset。dataset中的每一個元素就對應(yīng)了文件中的一行暴凑。可以使用這個函數(shù)來讀入CSV文件峦甩。
• tf.data.FixedLengthRecordDataset()：這個函數(shù)的輸入是一個文件的列表和一個record_bytes，之后dataset的每一個元素就是文件中固定字節(jié)數(shù)record_bytes的內(nèi)容。通常用來讀取以二進(jìn)制形式保存的文件凯傲，如CIFAR10數(shù)據(jù)集就是這種形式犬辰。
• tf.data.TFRecordDataset()：顧名思義，這個函數(shù)是用來讀TFRecord文件的冰单，dataset中的每一個元素就是一個TFExample幌缝。

它們的詳細(xì)使用方法可以參閱文檔：Module: tf.data

更多類型的Iterator....

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

• initializable iterator
• reinitializable iterator
• feedable iterator

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

limit = tf.placeholder(dtype=tf.int32, shape=[])

dataset = tf.data.Dataset.from_tensor_slices(tf.range(start=0, limit=limit))

iterator = dataset.make_initializable_iterator()
next_element = iterator.get_next()

with tf.Session() as sess:
    sess.run(iterator.initializer, feed_dict={limit: 10})
    for i in range(10):
      value = sess.run(next_element)
      assert i == value

此時(shí)的limit相當(dāng)于一個“參數(shù)”贴硫，它規(guī)定了Dataset中數(shù)的“上限”。

initializable iterator還有一個功能：讀入較大的數(shù)組伊者。

在使用tf.data.Dataset.from_tensor_slices(array)時(shí)，實(shí)際上發(fā)生的事情是將array作為一個tf.constants保存到了計(jì)算圖中间护。當(dāng)array很大時(shí)亦渗，會導(dǎo)致計(jì)算圖變得很大，給傳輸汁尺、保存帶來不便法精。這時(shí)，我們可以用一個placeholder取代這里的array痴突，并使用initializable iterator搂蜓，只在需要時(shí)將array傳進(jìn)去，這樣就可以避免把大數(shù)組保存在圖里辽装，示例代碼為（來自官方例程）：

# 從硬盤中讀入兩個Numpy數(shù)組
with np.load("/var/data/training_data.npy") as data:
  features = data["features"]
  labels = data["labels"]

features_placeholder = tf.placeholder(features.dtype, features.shape)
labels_placeholder = tf.placeholder(labels.dtype, labels.shape)

dataset = tf.data.Dataset.from_tensor_slices((features_placeholder, labels_placeholder))
iterator = dataset.make_initializable_iterator()
sess.run(iterator.initializer, feed_dict={features_placeholder: features,
                                          labels_placeholder: labels})

reinitializable iterator和feedable iterator相比initializable iterator更復(fù)雜帮碰，也更加少用，如果想要了解它們的功能拾积，可以參閱官方介紹殉挽，這里就不再贅述了。

總結(jié)

本文主要介紹了Dataset API的基本架構(gòu)：Dataset類和Iterator類拓巧，以及它們的基礎(chǔ)使用方法斯碌。

在非Eager模式下，Dataset中讀出的一個元素一般對應(yīng)一個batch的Tensor肛度，我們可以使用這個Tensor在計(jì)算圖中構(gòu)建模型傻唾。

在Eager模式下，Dataset建立Iterator的方式有所不同承耿，此時(shí)通過讀出的數(shù)據(jù)就是含有值的Tensor冠骄，方便調(diào)試伪煤。

作為兼容兩種模式的Dataset API，在今后應(yīng)該會成為TensorFlow讀取數(shù)據(jù)的主流方式猴抹。關(guān)于Dataset API的進(jìn)一步介紹带族，可以參閱下面的資料：

• Importing Data ：官方Guide
• Module: tf.data： API文檔
• Introduction to TensorFlow Datasets and Estimators：如何聯(lián)合使用Dataset和Estimator