1. 程式人生 > >TensorFlow全新的資料讀取方式:Dataset API入門教程

TensorFlow全新的資料讀取方式:Dataset API入門教程

Dataset API是TensorFlow 1.3版本中引入的一個新的模組,主要服務於資料讀取,構建輸入資料的pipeline。

此前,在TensorFlow中讀取資料一般有兩種方法:

相Dataset API同時支援從記憶體和硬碟的讀取,相比之前的兩種方法在語法上更加簡潔易懂。此外,如果想要用到TensorFlow新出的Eager模式,就必須要使用Dataset API來讀取資料。

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

Dataset API的匯入

在TensorFlow 1.3中,Dataset API是放在contrib包中的:

tf.contrib.data.Dataset

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

tf.data.Dataset

下面的示例程式碼將以TensorFlow 1.4版本為例,如果使用TensorFlow 1.3的話,需要進行簡單的修改(即加上contrib)。

基本概念:Dataset與Iterator

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

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

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

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

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]))

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

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

在非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))

對應的輸出結果應該就是從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異常,這個行為與使用佇列方式讀取資料的行為是一致的。在實際程式中,可以在外界捕捉這個異常以判斷資料是否讀取完,請參考下面的程式碼:

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模式中,建立Iterator的方式有所不同。是通過tfe.Iterator(dataset)的形式直接建立Iterator並迭代。迭代時可以直接取出值,不需要使用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)

從記憶體中建立更復雜的Dataset

之前我們用tf.data.Dataset.from_tensor_slices建立了一個最簡單的Dataset:

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

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

例如:

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

傳入的數值是一個矩陣,它的形狀為(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,例如我們可以讓每一個元素是一個詞典:

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))
    }
)

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

利用tf.data.Dataset.from_tensor_slices建立每個元素是一個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完成資料變換,打亂,組成batch,生成epoch等一系列操作。

常用的Transformation有:

  • map
  • batch
  • shuffle
  • repeat

下面就分別進行介紹。

(1)map

map接收一個函式,Dataset中的每個元素都會被當作這個函式的輸入,並將函式返回值作為新的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,如下面的程式將dataset中的每個元素組成了大小為32的batch:

dataset = dataset.batch(32)

(3)shuffle

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

dataset = dataset.shuffle(buffer_size=10000)

(4)repeat

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

dataset = dataset.repeat(5)

如果直接呼叫repeat()的話,生成的序列就會無限重複下去,沒有結束,因此也不會丟擲tf.errors.OutOfRangeError異常:

dataset = dataset.repeat()

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

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

對應的程式為(從官方示例程式修改而來):

# 函式的功能時將filename對應的圖片檔案讀進來,並縮放到統一的大小
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, ...])

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

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

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

在這個過程中,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。

更多型別的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 = 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

此時的limit相當於一個“引數”,它規定了Dataset中數的“上限”。

initializable iterator還有一個功能:讀入較大的陣列。

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

# 從硬碟中讀入兩個Numpy陣列
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更復雜,也更加少用,如果想要了解它們的功能,可以參閱官方介紹,這裡就不再贅述了。

總結

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

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

在Eager模式下,Dataset建立Iterator的方式有所不同,此時通過讀出的資料就是含有值的Tensor,方便除錯。

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