實例介紹TensorFlow的輸入流水線

520jefferson 2021-04-20

展開全文

作者：葉虎

編輯：趙一帆

前言

在訓(xùn)練模型時，我們首先要處理的就是訓(xùn)練數(shù)據(jù)的加載與預(yù)處理的問題，這里稱這個過程為輸入流水線（input pipelines，或輸入管道，[參考：https://www./performance/datasets_performance]）。在TensorFlow中，典型的輸入流水線包含三個流程（ETL流程）：

提取（Extract）：從存儲介質(zhì)（如硬盤）中讀取數(shù)據(jù)，可能是本地讀取，也可能是遠程讀?。ū热缭诜植际酱鎯ο到y(tǒng)HDFS）
預(yù)處理（Transform）：利用CPU處理器解析和預(yù)處理提取的數(shù)據(jù)，如圖像解壓縮，數(shù)據(jù)擴增或者變換，然后會做random shuffle，并形成batch。
加載（load）：將預(yù)處理后的數(shù)據(jù)加載到加速設(shè)備中（如GPUs）來執(zhí)行模型的訓(xùn)練。

輸入流水線對于加速模型訓(xùn)練還是很重要的，如果你的CPU處理數(shù)據(jù)能力跟不上GPU的處理速度，此時CPU預(yù)處理數(shù)據(jù)就成為了訓(xùn)練模型的瓶頸環(huán)節(jié)。除此之外，上述輸入流水線本身也有很多優(yōu)化的地方。比如，一個典型的模型訓(xùn)練過程中，CPU預(yù)處理數(shù)據(jù)時，GPU是閑置的，當(dāng)GPU訓(xùn)練模型時，CPU是閑置的，這個過程如下所示：

這樣一個訓(xùn)練step中所花費的時間是CPU預(yù)處理數(shù)據(jù)和GPU訓(xùn)練模型時間的總和。顯然這個過程中有資源浪費，一個改進的方法就是交叉CPU數(shù)據(jù)處理和GPU模型訓(xùn)練這兩個過程，當(dāng)GPU處于第個訓(xùn)練階段，CPU正在準(zhǔn)備第N+1步所需的數(shù)據(jù)，如下圖所示：

明顯上述設(shè)計可以充分最大化利用CPU和GPU，從而減少資源的閑置。另外當(dāng)存在多個CPU核心時，這又會涉及到CPU的并行化技術(shù)（多線程）來加速數(shù)據(jù)預(yù)處理過程，因為每個訓(xùn)練樣本的預(yù)處理過程往往是互相獨立的。關(guān)于輸入流程線的優(yōu)化可以參考TensorFlow官網(wǎng)上的Pipeline Performance Guide（https://www./performance/datasets_performance），相信你會受益匪淺。

幸運的是，最新的TensorFlow版本提供了tf.data這一套APIs來幫助我們快速實現(xiàn)高效又靈活的輸入流水線。在TensorFlow中最常見的加載訓(xùn)練數(shù)據(jù)的方式是通過Feeding（https://www./api_guides/python/reading_data#Feeding）方式，其主要是定義placeholder，然后將通過Session.run()的feed_dict參數(shù)送入數(shù)據(jù)，但是這其實是最低效的加載數(shù)據(jù)方式。后來，TensorFlow增加了QueueRunner（https://www./api_guides/python/reading_data#_QueueRunner_）機制，其主要是基于文件隊列以及多線程技術(shù)，實現(xiàn)了更高效的輸入流水線，但是其APIs很是讓人難懂，所以就有了現(xiàn)在的tf.data來替代它。

這里我們通過mnist實例來講解如何使用tf.data建立簡潔而高效的輸入流水線，在介紹之前，我們先介紹如何制作TFRecords文件，這是TensorFlow支持的一種標(biāo)準(zhǔn)文件格式

制作TFRecords文件

TFRecords文件是TensorFlow中的標(biāo)準(zhǔn)數(shù)據(jù)格式，它是基于protobuf的二進制文件，每個TFRecord文件的基本元素是tf.train.Example，其對應(yīng)的是數(shù)據(jù)集中的一個樣本數(shù)據(jù)，每個Example包含F(xiàn)eatures，存儲該樣本的各個feature，每個feature包含一個鍵值對，分別對應(yīng)feature的特征名與實際值。下面是一個Example實例：

// An Example for a movie recommendation application:
       features {
         feature {
           key: 'age'
           value { float_list {
             value: 29.0
           }}
         }
         feature {
           key: 'movie'
           value { bytes_list {
             value: 'The Shawshank Redemption'
             value: 'Fight Club'
           }}
         }
         feature {
           key: 'movie_ratings'
           value { float_list {
             value: 9.0
             value: 9.7
           }}
         }
         feature {
           key: 'suggestion'
           value { bytes_list {
             value: 'Inception'
           }}
         }
         feature {
           key: 'suggestion_purchased'
           value { float_list {
             value: 1.0
           }}
        }
         feature {
           key: 'purchase_price'
           value { float_list {
             value: 9.99
           }}
         }
      }

上面是一個電影推薦系統(tǒng)中的一個樣本，可以看到它共含有6個特征，每個特征都是key-value類型，key是特征名，而value是特征值，值得注意的是value其實存儲的是一個list，根據(jù)數(shù)據(jù)類型共分為三種：bytes_list, float_list和int64_list，分別存儲字節(jié)、浮點及整數(shù)類型（見這里：https://github.com/tensorflow/tensorflow/blob/r1.6/tensorflow/core/example/feature.proto）。

作為標(biāo)準(zhǔn)數(shù)據(jù)格式，TensorFlow當(dāng)然提供了創(chuàng)建TFRecords文件的python接口，下面我們創(chuàng)建mnist數(shù)據(jù)集對應(yīng)的TFRecords文件。對于mnist數(shù)據(jù)集，每個Example需要存儲兩個feature，一個是圖像的像素值，這里可以用bytes類型，因為一個像素點正好可以用一個字節(jié)存儲，另外是圖像的標(biāo)簽值，只能用int64類型存儲了。因此，我們先定義這兩個類型的接口函數(shù)：

 # int64
    def _int64_feature(value):
        return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))    # bytes
    def _bytes_feature(value):
        return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

創(chuàng)建TFRecord文件，主要通過TF中的tf.python_io.TFRecordWriter函數(shù)來實現(xiàn)，具體代碼如下：

def convert_to_TFRecords(dataset, name):
    '''Convert mnist dataset to TFRecords'''
    images, labels = dataset.images, dataset.labels
    n_examples = dataset.num_examples

    filename = os.path.join(DIR, name + '.tfrecords')
    print('Writing', filename)
    with tf.python_io.TFRecordWriter(filename) as writer:
            for index in range(n_examples):
            image_bytes = images[index].tostring()
            label = labels[index]
            example = tf.train.Example(features=tf.train.Features(
                feature={'image': _bytes_feature(image_bytes),                         'label': _int64_feature(label)}))
            writer.write(example.SerializeToString())

對于mnist數(shù)據(jù)集，主要分為train、validation和test，利用上面的函數(shù)分別創(chuàng)建三個不同的TFRecords文件：

mnist_datasets = mnist.read_data_sets('mnist_data', dtype=tf.uint8, reshape=False)
convert_to_TFRecords(mnist_datasets.train, 'train')
convert_to_TFRecords(mnist_datasets.validation, 'validation')
convert_to_TFRecords(mnist_datasets.test, 'test')

好了，這樣我們就創(chuàng)建3個TFRecords文件了。

讀取TFRecords文件

上面我們創(chuàng)建了TFRecords文件，但是怎么去讀取它們呢，當(dāng)然TF提供了讀取TFRecords文件的接口函數(shù)，這里首先介紹如何利用TF中操作TFRecord的python接口來讀取TFRecord文件，主要是tf.python_io.tf_record_iterator函數(shù)，它輸入TFRecord文件，但是得到一個迭代器，每個元素是一個Example，但是卻是一個字符串，這里可以用tf.train.Example來解析它，具體代碼如下：

def read_TFRecords_test(name):
    filename = os.path.join(DIR, name + '.tfrecords')
    record_itr = tf.python_io.tf_record_iterator(path=filename)
    for r in record_itr:
        example = tf.train.Example()
        example.ParseFromString(r)

        label = example.features.feature['label'].int64_list.value[0]
        print('Label', label)
        image_bytes = example.features.feature['image'].bytes_list.value[0]
        img = np.fromstring(image_bytes, dtype=np.uint8).reshape(28, 28)
        print(img)
        plt.imshow(img, cmap='gray')
        plt.show()
        break  # 只讀取一個Example

上面僅是純python的讀取方式，這不是TFRecords文件的正確使用方式。既然是官方標(biāo)準(zhǔn)數(shù)據(jù)格式，TF也提供了使用TFRecords文件建立輸入流水線的方式。在tf.data出現(xiàn)之前，使用的是QueueRunner方式，即文件隊列機制，其原理如下圖所示：

文件隊列機制主要分為兩個階段：第一個階段將輸入文件打亂，并在文件隊列入列，然后Reader從文件隊列中讀取一個文件，同時文件隊列出列這個文件，Reader同時對文件進行解碼，然后生產(chǎn)數(shù)據(jù)樣本，并將樣本在樣本隊列中入列，可以定義多個Reader并發(fā)地從多個文件同時讀取數(shù)據(jù)。從樣本隊列中的出列一定量的樣本數(shù)據(jù)即可以用于一個訓(xùn)練過程。TF提供了配套的API來完成這個過程，注意的是這個輸入流水線是直接嵌入訓(xùn)練的Graph中，即是整個圖模型的一部分。根據(jù)文件的不同，可以使用不同類型的Reader，對于TFRecord文件，可以使用tf.TFRecordReader，下面是具體的實現(xiàn)代碼：

def read_example(filename_queue):
    '''Read one example from filename_queue'''
    reader = tf.TFRecordReader()
    key, value = reader.read(filename_queue)
    features = tf.parse_single_example(value, features={'image': tf.FixedLenFeature([], tf.string),                                                            'label': tf.FixedLenFeature([], tf.int64)})
    image = tf.decode_raw(features['image'], tf.uint8)
    image = tf.reshape(image, [28, 28])
    label = tf.cast(features['label'], tf.int32)
    return image, label
    
if __name__ == '__main__':
    queue = tf.train.string_input_producer(['TFRecords/train.tfrecords'], num_epochs=10)
    image, label = read_example(queue)

    img_batch, label_batch = tf.train.shuffle_batch([image, label], batch_size=32, capacity=5000,
                                                        min_after_dequeue=2000, num_threads=4)
    with tf.Session() as sess:
        sess.run(tf.local_variables_initializer())
        sess.run(tf.global_variables_initializer())

        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)
        try:
                while not coord.should_stop():                # Run training steps or whatever
                images, labels = sess.run([img_batch, label_batch])
                print(images.shape, labels.shape)
        except tf.errors.OutOfRangeError:
            print('Done training -- epoch limit reached')

        coord.request_stop()
        coord.join(threads)

對于隊列機制，估計大家看的云里霧里的，代碼確實讓人難懂，但是其實只要按照官方提供的標(biāo)準(zhǔn)代碼，還是很容易在自己的數(shù)據(jù)集上進行修改的。不過現(xiàn)在有了tf.data，可以更加優(yōu)雅地實現(xiàn)上面的過程。

tf.data簡介

使用tf.data可以更方便地創(chuàng)建高效的輸入流水線，但是其相比隊列機制API更友好，這主要是因為tf.data提供了高級抽象。第一個抽象是使用tf.data.Dataset來表示一個數(shù)據(jù)集合，集合里面的每個元素包含一個或者多個Tensor，一般就是對應(yīng)一個訓(xùn)練樣本。第二個抽象是使用tf.data.Iterator來從數(shù)據(jù)集中提取數(shù)據(jù)，這是一個迭代器對象，可以通過Iterator.get_next()從Dataset中產(chǎn)生一個樣本。利用這兩個抽象，Dataset的使用簡化為三個步驟：

創(chuàng)建Dataset實例對象；
創(chuàng)建遍歷Dataset的Iterator實例對象；
從Iterator中不斷地產(chǎn)生樣本，并送入模型中進行訓(xùn)練。

1.創(chuàng)建Dataset

TF提供了很多方式創(chuàng)建Dataset，下面是幾種方式：

# 從Numpy的arraydataset1 = tf.data.Dataset.from_tensor_slices(np.random.randn((5, 10))
print(dataset1.output_types)  # ==> 'tf.float32'print(dataset1.output_shapes)  # ==> '(10,)'# 從Tensor

dataset2 = tf.data.Dataset.from_tensor_slices((tf.random_uniform([4]),
tf.random_uniform([4, 100], maxval=100, dtype=tf.int32)))
print(dataset2.output_types)  # ==> '(tf.float32, tf.int32)'print(dataset2.output_shapes)  # ==> '((), (100,))'# 從文件
filenames = ['/var/data/file1.tfrecord', '/var/data/file2.tfrecord']
dataset3 = tf.data.TFRecordDataset(filenames)

更重要的是Dataset可以進行一系列的變換操作，并且支持鏈?zhǔn)秸{(diào)用，這對于數(shù)據(jù)預(yù)處理很重要：

dataset = tf.data.TFRecordDataset(filenames)
dataset = dataset.map(...)  # 解析數(shù)據(jù)或者對數(shù)據(jù)預(yù)處理，如normalize.

dataset = dataset.repeat()  # 重復(fù)數(shù)據(jù)集，一般設(shè)置num_epochs

dataset = dataset.batch(32) # 形成batch

2.創(chuàng)建Iterator

創(chuàng)建了Dataset之后，我們需要創(chuàng)建Iterator來遍歷數(shù)據(jù)集，返回的是迭代器對象，并從中可以產(chǎn)生數(shù)據(jù)，以用于模型訓(xùn)練。TF共支持4中迭代器類型，分別是one-shot, initializable, reinitializable和feedable。下面逐個介紹它們。

One-shot Iterator

這是最簡單的Iterator，它僅僅遍歷整個數(shù)據(jù)集一次，而且不需要顯示初始化，下面是個實例：

dataset = tf.data.Dataset.from_tensor_slices(np.arange(10))
iterator = dataset.make_one_shot_iterator()
next_element = iterator.get_next()
with tf.Session() as sess:
    for i in range(10):
        sess.run(next_element) # 0, 1, ..., 9

Initializable Iterator

相比one-shot Iterator，它需要在使用前顯示初始化，這樣就可以支持參數(shù)化，每次初始化時送入不同的參數(shù)，就可以支持?jǐn)?shù)據(jù)集的簡單參數(shù)化，下面是一個實例：

max_value = tf.placeholder(tf.int64, [])
dataset = tf.data.Dataset.range(max_value)
iterator = dataset.make_initializable_iterator()
next_element = iterator.get_next()
with tf.Session() as sess:    # 需要顯示初始化
    sess.run(iterator.initializer, feed_dict={max_value: 10})
    for i in range(10):
        print(sess.run(next_element)) # 0, 1, ..., 9

Reinitializable Iterator

相比initializable Iterator，它可以支持從不同的Dataset進行初始化，有時候你需要訓(xùn)練集和測試集，但是兩者并不同，此時就可以定義兩個不同的Dataset，并配合reinitializable Iterator來定義一個通用的迭代器，在使用前只需要送入不同的Dataset進行初始化就可以，下面是一個實例：

train_data = np.random.randn(100, 5)
test_data = np.random.randn(20, 5)
train_dataset = tf.data.Dataset.from_tensor_slices(train_data)
test_dataset = tf.data.Dataset.from_tensor_slices(test_data)# 創(chuàng)建一個reinitializable iterator

re_iterator = tf.data.Iterator.from_structure(train_dataset.output_types,
                                              train_dataset.output_shapes)
next_element = re_iterator.get_next()
train_init_op = re_iterator.make_initializer(train_dataset)
test_init_op = re_iterator.make_initializer(test_dataset)
with tf.Session() as sess:    # 訓(xùn)練
    n_epochs = 2
    for i in range(n_epochs):
        sess.run(train_init_op)
        for j in range(100):
            print(sess.run(next_element))
    # 測試
    sess.run(test_init_op)
    for i in range(20):
        print(sess.run(next_element))

Feedable Iterator

對于reinitializable iterator，它可以支持送入不同Dataset，從而完成數(shù)據(jù)集的切換，但是每次切換時必須要重新初始化。對于Feedable Iterator，其可以認為支持送入不同的Iterator，通過切換迭代器的string handle來完成不同數(shù)據(jù)集的切換，并且在切換時迭代器的狀態(tài)還會被保留，這相比reinitializable iterator更加靈活，下面是一個實例：

train_data = np.random.randn(100, 5)
val_data = np.random.randn(20, 5)
n_epochs = 20train_dataset = tf.data.Dataset.from_tensor_slices(train_data).repeat(n_epochs)
val_dataset = tf.data.Dataset.from_tensor_slices(val_data)# 創(chuàng)建一個feedable iterator

handle = tf.placeholder(tf.string, [])
feed_iterator = tf.data.Iterator.from_string_handle(handle, train_dataset.output_types,
                                                  train_dataset.output_shapes)
next_element = feed_iterator.get_next()# 創(chuàng)建不同的iterator

train_iterator = train_dataset.make_one_shot_iterator()
val_iterator = val_dataset.make_initializable_iterator()
with tf.Session() as sess:
    # 生成對應(yīng)的handle
    train_handle = sess.run(train_iterator.string_handle())
    val_handle = sess.run(val_iterator.string_handle())
    # 訓(xùn)練
    for n in range(n_epochs):
        for i in range(100):
            print(i, sess.run(next_element, feed_dict={handle: train_handle}))
                # 驗證
        if n % 10 == 0:
            sess.run(val_iterator.initializer)
                    for i in range(20):
                print(sess.run(next_element, feed_dict={handle: val_handle}))

關(guān)于tf.data的基礎(chǔ)知識就這么多了，更多內(nèi)容可以參考官方文檔，另外這里要說一點就是，對于迭代器對象，當(dāng)其元素取盡之后，會拋出tf.errors.OutOfRangeError錯誤，當(dāng)然一般情況下你是知道自己的迭代器對象的元素數(shù)，那么也就可以不用通過捕獲錯誤來實現(xiàn)終止條件。下面，我們將使用tf.data實現(xiàn)mnist的完整訓(xùn)練過程。

MNIST完整實例

我們采用feedable Iterator來實現(xiàn)mnist數(shù)據(jù)集的訓(xùn)練過程，分別創(chuàng)建兩個Dataset，一個為訓(xùn)練集，一個為驗證集，對于驗證集不需要shuffle操作。首先我們創(chuàng)建Dataset對象的輔助函數(shù)，主要是解析TFRecords文件，并對image做歸一化處理：

def decode(serialized_example):
    '''decode the serialized example'''
    features = tf.parse_single_example(serialized_example,
                            features={'image': tf.FixedLenFeature([], tf.string),                                      'label': tf.FixedLenFeature([], tf.int64)})
    image = tf.decode_raw(features['image'], tf.uint8)
    image = tf.cast(image, tf.float32)
    image = tf.reshape(image, [784])
    label = tf.cast(features['label'], tf.int64)
    return image, label
def normalize(image, label):
    '''normalize the image to [-0.5, 0.5]'''
    image = image / 255.0 - 0.5
    return image, label

然后定義創(chuàng)建Dataset的函數(shù)，對于訓(xùn)練集和驗證集，兩者的參數(shù)會不同：

def create_dataset(filename, batch_size=64, is_shuffle=False, n_repeats=0):
    '''create dataset for train and validation dataset'''
    dataset = tf.data.TFRecordDataset(filename)
    if n_repeats > 0:
        dataset = dataset.repeat(n_repeats) # for train

    dataset = dataset.map(decode).map(normalize) # decode and normalize

    if is_shuffle:
        dataset = dataset.shuffle(1000 + 3 * batch_size) # shuffle
    dataset = dataset.batch(batch_size)
    return dataset

我們使用一個簡單的全連接層網(wǎng)絡(luò)來實現(xiàn)mnist的分類模型：

def model(inputs, hidden_sizes=(500, 500)):
    h1, h2 = hidden_sizes
    net = tf.layers.dense(inputs, h1, activation=tf.nn.relu)
    net = tf.layers.dense(net, h2, activation=tf.nn.relu)
    net = tf.layers.dense(net, 10, activation=None)
    return net

然后是訓(xùn)練的主體代碼：

n_train_examples = 55000n_val_examples = 5000n_epochs = 50batch_size = 64train_dataset = create_dataset('TFRecords/train.tfrecords', batch_size=batch_size, is_shuffle=True,
                               n_repeats=n_epochs)
val_dataset = create_dataset('TFRecords/validation.tfrecords', batch_size=batch_size)
# 創(chuàng)建一個feedable iterator

handle = tf.placeholder(tf.string, [])
feed_iterator = tf.data.Iterator.from_string_handle(handle, train_dataset.output_types,
                                                  train_dataset.output_shapes)
images, labels = feed_iterator.get_next()
# 創(chuàng)建不同的iterator

train_iterator = train_dataset.make_one_shot_iterator()
val_iterator = val_dataset.make_initializable_iterator()
# 創(chuàng)建模型

logits = model(images, [500, 500])
loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels, logits=logits)
loss = tf.reduce_mean(loss)
train_op = tf.train.AdamOptimizer(learning_rate=1e-04).minimize(loss)
predictions = tf.argmax(logits, axis=1)
accuracy = tf.reduce_mean(tf.cast(tf.equal(predictions, labels), tf.float32))

init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
with tf.Session() as sess:
    sess.run(init_op)    # 生成對應(yīng)的handle
    train_handle = sess.run(train_iterator.string_handle())
    val_handle = sess.run(val_iterator.string_handle())
    # 訓(xùn)練
    for n in range(n_epochs):
        ls = []
        for i in range(n_train_examples // batch_size):
            _, l = sess.run([train_op, loss], feed_dict={handle: train_handle})
            ls.append(l)
        print('Epoch %d, train loss: %f' % (n, np.mean(ls)))
        if (n + 1) % 10 == 0:
            sess.run(val_iterator.initializer)
            accs = []
            for i in range(n_val_examples // batch_size):
                acc = sess.run(accuracy, feed_dict={handle: val_handle})
                accs.append(acc)
            print('\t validation accuracy: %f' % (np.mean(accs)))

大約可以在驗證集上的accuracy達到98%。

小結(jié)

看起來最新的tf.data還是比較好用的，如果你是TensorFlow用戶，可以嘗試著使用它，當(dāng)然上面的例子并不能包含關(guān)于tf.data的所有內(nèi)容，想繼續(xù)深入的話可以移步TF的官網(wǎng)。

參考資料

[Programmers guide: import data](https://www./programmers_guide/datasets).
[How to use Dataset in TensorFlow](https:///how-to-use-dataset-in-tensorflow-c758ef9e4428).
[Reading data](https://www./api_guides/python/reading_data).
[Performance: datasets performance](https://www./performance/datasets_performance).
[Introduction to Artificial Neural Networks and Deep Learning: A Practical Guide with Applications in Python](https://github.com/rasbt/deep-learning-book/).

end