GPU 和 TPU 可以從根本上縮短執(zhí)行單個(gè)訓(xùn)練步驟所需的時(shí)間。欲將性能提高到極致，則需要有一個(gè)高效的輸入管道，能夠在當(dāng)前步驟完成之前為下一步提供數(shù)據(jù)。tf.data API 有助于構(gòu)建靈活高效的輸入管道。本文介紹了 tf.data API 的功能和最佳實(shí)踐操作，用于在各種模型和加速器上構(gòu)建高性能 TensorFlow 輸入管道。

本文將主要介紹下列內(nèi)容：

說明 TensorFlow 輸入管道本質(zhì)上是一個(gè)ETL過程

描述 tf.dataAPI 上下文中的常見性能優(yōu)化

討論應(yīng)用轉(zhuǎn)換順序的性能影響

總結(jié)設(shè)計(jì)高性能 TensorFlow 輸入管道的最佳實(shí)踐操作。

輸入數(shù)據(jù)管道的結(jié)構(gòu)

一個(gè)典型的 TensorFlow 訓(xùn)練輸入管道可以構(gòu)建為 ETL 過程：

提取：從持久存儲(chǔ)中讀取數(shù)據(jù) - 本地（例如 HDD 或 SSD）或遠(yuǎn)程（例如 GCS 或 HDFS）

轉(zhuǎn)換：使用 CPU 內(nèi)核對數(shù)據(jù)進(jìn)行解析和執(zhí)行預(yù)處理操作，例如圖像解壓縮，數(shù)據(jù)擴(kuò)充轉(zhuǎn)換（例如隨機(jī)裁剪，翻轉(zhuǎn)和顏色失真），隨機(jī)洗牌和批處理

加載：將變換后的數(shù)據(jù)加載到執(zhí)行機(jī)器學(xué)習(xí)模型的加速器設(shè)備（例如，GPU 或 TPU）上。

這種模式有效地利用了 CPU，與此同時(shí)為了模型訓(xùn)練這種繁重的工作，還保留了加速器。此外，將輸入管道視為 ETL 過程提供了便于性能優(yōu)化應(yīng)用的結(jié)構(gòu)。

使用 tf.estimator.EstimatorAPI 時(shí)，前兩個(gè)階段（提取和轉(zhuǎn)換）將在傳遞給 tf.estimator.Estimator 的 input_fn 中捕獲。在代碼中，這可能會(huì)看起來像以下（naive, sequential）執(zhí)行情況：

defparse_fn(example): "Parse TFExample records and perform simple data augmentation." example_fmt = { "image": tf.FixedLengthFeature((), tf.string, ""), "label": tf.FixedLengthFeature((), tf.int64, -1) } parsed = tf.parse_single_example(example, example_fmt) image = tf.image.decode_image(parsed["image"]) image = _augment_helper(image) # augments image using slice, reshape, resize_bilinear returnimage, parsed["label"]definput_fn(): files = tf.data.Dataset.list_files("/path/to/dataset/train-*.tfrecord") dataset = files.interleave(tf.data.TFRecordDataset) dataset = dataset.shuffle(buffer_size=FLAGS.shuffle_buffer_size) dataset = dataset.map(map_func=parse_fn) dataset = dataset.batch(batch_size=FLAGS.batch_size) returndataset

下一部分將基于此輸入管道構(gòu)建，添加性能優(yōu)化。

性能優(yōu)化

首先，我們定義要使用的模型種類。主智能體會(huì)擁有全局網(wǎng)絡(luò)，且每個(gè)本地工作器智能體在自己的進(jìn)程中都會(huì)擁有此網(wǎng)絡(luò)的副本。我們會(huì)使用模型子類化對模型進(jìn)行實(shí)例化。雖然模型子類化會(huì)使進(jìn)程更冗長，但卻為我們提供了最大的靈活性。

正如您從我們的正向傳遞中看到的，我們的模型會(huì)采用輸入和返回策略概率的分對數(shù)和值。

隨著新的計(jì)算設(shè)備（諸如 GPU 和 TPU）不斷問世，訓(xùn)練神經(jīng)網(wǎng)絡(luò)的速度變得越來越快，這種情況下 CPU 處理很容易成為瓶頸。tf.dataAPI 為用戶提供構(gòu)建塊，以設(shè)計(jì)有效利用 CPU 的輸入管道，優(yōu)化 ETL 過程的每個(gè)步驟。

Pipelining

要執(zhí)行訓(xùn)練步驟，您必須首先提取并轉(zhuǎn)換訓(xùn)練數(shù)據(jù)，然后將其提供給在加速器上運(yùn)行的模型。然而，在一個(gè)簡單的同步執(zhí)行中，當(dāng) CPU 正在準(zhǔn)備數(shù)據(jù)時(shí)，加速器則處于空閑狀態(tài)。相反，當(dāng)加速器正在訓(xùn)練模型時(shí)，CPU 則處于空閑狀態(tài)。因此，訓(xùn)練步驟時(shí)間是 CPU 預(yù)處理時(shí)間和加速器訓(xùn)練時(shí)間的總和。

Pipelining 將一個(gè)訓(xùn)練步驟的預(yù)處理和模型執(zhí)行重疊。當(dāng)加速器正在執(zhí)行訓(xùn)練步驟 N 時(shí)，CPU 正在準(zhǔn)備步驟 N + 1 的數(shù)據(jù)。這樣做的目的是可以將步驟時(shí)間縮短到極致，包含訓(xùn)練以及提取和轉(zhuǎn)換數(shù)據(jù)所需時(shí)間（而不是總和）。

如果沒有使用 pipelining，則 CPU 和 GPU / TPU 在大部分時(shí)間處于閑置狀態(tài)：

而使用 pipelining 技術(shù)后，空閑時(shí)間顯著減少：

tf.dataAPI 通過 tf.data.Dataset.prefetch 轉(zhuǎn)換提供了一個(gè)軟件 pipelining 操作機(jī)制，該轉(zhuǎn)換可用于將數(shù)據(jù)生成的時(shí)間與所消耗時(shí)間分離。特別是，轉(zhuǎn)換使用后臺線程和內(nèi)部緩沖區(qū)，以便在請求輸入數(shù)據(jù)集之前從輸入數(shù)據(jù)集中預(yù)提取元素。因此，為了實(shí)現(xiàn)上面說明的 pipelining 效果，您可以將 prefetch(1) 添加為數(shù)據(jù)集管道的最終轉(zhuǎn)換（如果單個(gè)訓(xùn)練步驟消耗 n 個(gè)元素，則添加 prefetch(n)）。

要將此更改應(yīng)用于我們的運(yùn)行示例，請將：

dataset = dataset.batch(batch_size=FLAGS.batch_size)returndataset

更改為：

dataset = dataset.batch(batch_size=FLAGS.batch_size)dataset = dataset.prefetch(buffer_size=FLAGS.prefetch_buffer_size)returndataset

請注意，在任何時(shí)候只要有機(jī)會(huì)將 “制造者” 的工作與 “消費(fèi)者” 的工作重疊，預(yù)取轉(zhuǎn)換就會(huì)產(chǎn)生效益。前面的建議只是最常見的應(yīng)用程序。

將數(shù)據(jù)轉(zhuǎn)換并行化

準(zhǔn)備批處理時(shí)，可能需要預(yù)處理輸入元素。為此，tf.dataAPI 提供了 tf.data.Dataset.map 轉(zhuǎn)換，它將用戶定義的函數(shù)（例如，運(yùn)行示例中的 parse_fn）應(yīng)用于輸入數(shù)據(jù)集的每個(gè)元素。由于輸入元素彼此獨(dú)立，因此可以跨多個(gè) CPU 內(nèi)核并行化預(yù)處理。為了實(shí)現(xiàn)這一點(diǎn)，map 轉(zhuǎn)換提供了 thenum_parallel_calls 參數(shù)來指定并行度。例如，下圖說明了將 num_parallel_calls = 2 設(shè)置為 map 轉(zhuǎn)換的效果：

為 num_parallel_calls 參數(shù)選擇最佳值取決于您的硬件，訓(xùn)練數(shù)據(jù)的特征（例如其大小和形狀），Map 功能的成本以及在 CPU 上同時(shí)進(jìn)行的其他處理；一個(gè)簡單的啟發(fā)式方法是使用可用的 CPU 內(nèi)核數(shù)。例如，如果執(zhí)行上述示例的機(jī)器有 4 個(gè)內(nèi)核，則設(shè)置 num_parallel_calls = 4 會(huì)更有效。另一方面，將 num_parallel_calls 設(shè)置為遠(yuǎn)大于可用 CPU 數(shù)量的值可能會(huì)導(dǎo)致調(diào)度效率低下，從而導(dǎo)致速度減慢。

要將此更改應(yīng)用于我們的運(yùn)行示例，請將：

dataset = dataset.map(map_func=parse_fn)

變更為：

dataset = dataset.map(map_func=parse_fn, num_parallel_calls=FLAGS.num_parallel_calls)

此外，如果您的批處理大小為數(shù)百或數(shù)千，您的 pipeline 可能還可以通過并行化批處理創(chuàng)建而從中獲益。為此，tf.dataAPI 提供了 tf.contrib.data.map_and_batch 轉(zhuǎn)換，它有效地 “融合” 了 map 和批處理的轉(zhuǎn)換。

要將此更改應(yīng)用于我們的運(yùn)行示例，請將：

dataset = dataset.map(map_func=parse_fn, num_parallel_calls=FLAGS.num_parallel_calls)dataset = dataset.batch(batch_size=FLAGS.batch_size)

更改為:

dataset = dataset.apply(tf.contrib.data.map_and_batch( map_func=parse_fn, batch_size=FLAGS.batch_size))

將數(shù)據(jù)提取并行化

在實(shí)際環(huán)境中，輸入數(shù)據(jù)可能被遠(yuǎn)程存儲(chǔ)（例如，GCS 或 HDFS），因?yàn)檩斎霐?shù)據(jù)不適合本地，或者因?yàn)橛?xùn)練是分布式的，因此在每臺機(jī)器上復(fù)制輸入數(shù)據(jù)是沒有意義的。在本地讀取數(shù)據(jù)時(shí)運(yùn)行良好的數(shù)據(jù)集管道在遠(yuǎn)程讀取數(shù)據(jù)時(shí)可能會(huì)成為 I / O 的瓶頸，因?yàn)楸镜卮鎯?chǔ)和遠(yuǎn)程存儲(chǔ)之間存在以下差異：

首字節(jié)時(shí)間：從遠(yuǎn)程存儲(chǔ)中讀取文件的第一個(gè)字節(jié)可能比本地存儲(chǔ)長幾個(gè)數(shù)量級

讀取吞吐量：雖然遠(yuǎn)程存儲(chǔ)通常提供較大的聚合帶寬，但讀取單個(gè)文件可能只能使用此帶寬的一小部分。

另外，一旦將原始字節(jié)讀入存儲(chǔ)器，也可能需要對數(shù)據(jù)進(jìn)行反序列化或解密（例如，protobuf），這就增加了額外的系統(tǒng)開銷。無論數(shù)據(jù)是本地存儲(chǔ)還是遠(yuǎn)程存儲(chǔ)，都存在這種開銷，如果數(shù)據(jù)未被有效預(yù)取，則在遠(yuǎn)程情況下情況可能更糟。

為了減輕各種數(shù)據(jù)提取開銷的影響，tf.dataAPI 提供了 tf.contrib.data.parallel_interleave 轉(zhuǎn)換。使用此轉(zhuǎn)換可以將其他數(shù)據(jù)集（例如數(shù)據(jù)文件讀取器）的內(nèi)容執(zhí)行和交錯(cuò)并行化。可以通過 cycle_length 參數(shù)指定要重疊的數(shù)據(jù)集的數(shù)量。

為 parallel_interleavetransformation 提供 cycle_length = 2 的效果如下圖所示：

要將此更改應(yīng)用于我們的運(yùn)行示例，請將：

dataset = files.interleave(tf.data.TFRecordDataset)

更改為:

dataset = files.apply(tf.contrib.data.parallel_interleave( tf.data.TFRecordDataset, cycle_length=FLAGS.num_parallel_readers))

由于負(fù)載或網(wǎng)絡(luò)事件，遠(yuǎn)程存儲(chǔ)系統(tǒng)的吞吐量可能會(huì)隨時(shí)間而變化。為了解釋這種差異，parallel_interleave 轉(zhuǎn)換可以選擇使用預(yù)取。（請參考 tf.contrib.data.parallel_interleave 了解詳情 https://tensorflow.google.cn/api_docs/python/tf/contrib/data/parallel_interleave?hl=zh-CN）。

默認(rèn)情況下，parallel_interleave 轉(zhuǎn)換提供了元素的確定性排序使之重現(xiàn)。作為預(yù)取的替代方法（在某些情況下可能無效），parallel_interleave 轉(zhuǎn)換還提供了一個(gè)選項(xiàng)，能夠以保證排序作為代價(jià)來提高性能。尤其是如果 sloppy 參數(shù)設(shè)置為 true，則轉(zhuǎn)換可能會(huì)偏離其確定的排序，在請求下一元素時(shí)那些不可用文件將會(huì)暫時(shí)跳過。

性能注意事項(xiàng)

tf.dataAPI 圍繞可組合轉(zhuǎn)換而設(shè)計(jì)，為用戶提供了靈活性。雖然這些轉(zhuǎn)換中的許多都是可交換的，但某些轉(zhuǎn)換的排序具有性能上的影響。

Map 和 batch

調(diào)用傳遞給 map 轉(zhuǎn)換的用戶定義函數(shù)會(huì)帶來與調(diào)度和執(zhí)行用戶定義函數(shù)相關(guān)的系統(tǒng)開銷。通常，與函數(shù)執(zhí)行的計(jì)算量相比，這種系統(tǒng)開銷很小。但是，如果 map 幾乎沒有使用，那么這種開銷可能會(huì)占據(jù)總成本的大多數(shù)。在這種情況下，我們建議對用戶定義的函數(shù)進(jìn)行矢量化（即，讓它一次對一批輸入進(jìn)行操作），并在 map 轉(zhuǎn)換之前應(yīng)用 batch 轉(zhuǎn)換。

Map 和 cache

tf.data.Dataset.cache 轉(zhuǎn)換可以在內(nèi)存或本地存儲(chǔ)中緩存數(shù)據(jù)集。如果傳遞給 map 轉(zhuǎn)換的用戶定義函數(shù)非常高，只要結(jié)果數(shù)據(jù)集仍然適合內(nèi)存或本地存儲(chǔ)，就可以在 map 轉(zhuǎn)換后應(yīng)用緩存轉(zhuǎn)換。如果用戶定義的函數(shù)增加了存儲(chǔ)數(shù)據(jù)集所需的空間超出緩存容量，請考慮在訓(xùn)練作業(yè)之前預(yù)處理數(shù)據(jù)以減少資源使用。

Map 和Interleave / Prefetch / Shuffle

許多轉(zhuǎn)換（包括 Interleave，prefetch 和 shuffle）會(huì)保留元素的內(nèi)部緩沖區(qū)。如果傳遞給 map 變換的用戶定義函數(shù)改變了元素的大小，那么 map 變換的順序和緩沖元素的變換會(huì)影響內(nèi)存使用。通常來說，除非由于性能需要不同的排序（例如，啟用 map 和 batch 轉(zhuǎn)換的融合）的情況，否則我們建議選擇帶來較低內(nèi)存占用的順序。

Repeat 和 Shuffle

tf.data.Dataset.repeat 轉(zhuǎn)換以有限（或無限）次數(shù)重復(fù)輸入數(shù)據(jù); 每次數(shù)據(jù)重復(fù)通常稱為 epoch。

tf.data.Dataset.shuffle 轉(zhuǎn)換隨機(jī)化數(shù)據(jù)集示例的順序。

如果在 shuffle 變換之前應(yīng)用 repeat 變換，則 epoch 的邊界模糊。也就是說，某些元素可以在其他元素出現(xiàn)之前重復(fù)一次。另一方面，如果在 repeat 變換之前應(yīng)用 shuffle 變換，則性能可能在與 shuffle 轉(zhuǎn)換的內(nèi)部狀態(tài)的初始化相關(guān)的每個(gè) epoch 時(shí)期的開始時(shí)減慢。換句話說，前者（repeat before shuffle）提供更好的性能，而后者（shuffle before repeat）提供更強(qiáng)的排序保證。

如果可能，我們推薦使用融合的 tf.contrib.data.shuffle_and_repeat 轉(zhuǎn)換，它結(jié)合了兩方面的優(yōu)點(diǎn)（良好的性能和強(qiáng)大的排序保證）。否則，我們建議在 repeating 之前進(jìn)行 shuffling。

最佳的實(shí)踐操作摘要

以下是設(shè)計(jì)輸入管道的最佳實(shí)踐操作摘要：

使用 prefetch 轉(zhuǎn)換重疊 “制造者” 和 “消費(fèi)者” 的工作。特別是，我們建議將 prefetch（n）（其中 n 是訓(xùn)練步驟消耗的元素 / 批次數(shù)）添加到輸入管道的末尾，以便在 CPU 上執(zhí)行的轉(zhuǎn)換與加速器上的訓(xùn)練重疊

通過設(shè)置 num_parallel_calls 參數(shù)來并行化 map 轉(zhuǎn)換。我們建議使用可用 CPU 內(nèi)核數(shù)作為其參數(shù)值

如果使用 batch 轉(zhuǎn)換將預(yù)處理元素組合成批處理，我們建議使用融合的 map_and_batch 轉(zhuǎn)換，特別是在您使用大型批處理的情況下

如果您正在處理遠(yuǎn)程存儲(chǔ)的數(shù)據(jù)和 / 或需要反序列化，我們建議使用 parallel_interleave 轉(zhuǎn)換來重疊來自不同文件的數(shù)據(jù)的讀取（和反序列化）

向傳遞到 map 轉(zhuǎn)換的廉價(jià)用戶定義函數(shù)進(jìn)行向量化，以分?jǐn)偱c調(diào)度和執(zhí)行函數(shù)相關(guān)的系統(tǒng)開銷

如果您的數(shù)據(jù)可以存儲(chǔ)于內(nèi)存中，請使用 cache 轉(zhuǎn)換在第一個(gè) epoch 期間將其緩存在內(nèi)存中，以便后續(xù) epoch 期間避免發(fā)生與讀取，解析和轉(zhuǎn)換相關(guān)的系統(tǒng)開銷

如果預(yù)處理增加了數(shù)據(jù)的大小，我們建議您首先應(yīng)用 interleave，prefetch 和 shuffle（如果可能的話）以減少內(nèi)存使用量

我們建議在 repeat 轉(zhuǎn)換之前應(yīng)用 shuffle 轉(zhuǎn)換，理想情況下使用融合的 shuffle_and_repeat 轉(zhuǎn)換。