TensorFlow 高效能資料輸入管道設計指南
作者：黑暗星球
原文地址：https://blog.csdn.net/u014061630/article/details/80776975
TensorFlow版本：1.12.0
本篇主要介紹怎麼使用 tf.data API 來構建高效能的輸入 pipeline。
tf.data官方教程詳見前面的部落格<<<<<<<<<<tf.data官方教程
文章目錄
- 1. 資料輸入管道的結構
- 2. 資料輸入管道的效能優化
- 2.1 資料準備(ET)和消耗過程(L)的解耦 ------
prefetch解耦、重疊兩過程 - 2.2 資料變換(T)的並行化 ------ 並行
map，融合map和batch - 2.3 資料讀取(E)的並行化------並行地讀取並解析多個數據檔案
- 3. 資料輸入管道效能的進一步優化
- 3.1 Map and Batch ------ map開銷很小時，batch形式的map更高效
- 3.2 Map and Cache ------ 通過cache進一步加速
- 3.3 Map and Interleave / Prefetch / Shuffle ------ 變換的順序對記憶體使用量的影響
- 3.4 Repeat and Shuffle ------ repeat在前效能優，shuffle在前次序強
- 4. 資料輸入管道的最優實現
GPU、TPU的使用能夠從根本上減少單個訓練step所需的時間。但優異的效能不僅依賴於高速的計算硬體，也要求有一個高效的輸入管道（Input Pipeline Performance Guide），這個管道在當前step完成前，進行下一個 step 需要的資料的準備。
tf.data API 對於靈活且高效的輸入管道的建立非常有幫助。這個文件解釋了
tf.data API 的特性，並介紹了構建高效能的 TensorFlow 資料輸入管道的過程。
本文主要包含以下內容：
- 介紹資料輸入管道的結構（本質是一個 ETL 過程）。
- 在
tf.data中，優化資料輸入管道的常用方法。 - 介紹了資料操作順序對資料輸入管道效能的影響。
- 優異的資料輸入管道應該具備的一些特質。
1. 資料輸入管道的結構
TensorFlow資料輸入管道可以被抽象為一個 ETL 過程（Extract，Transform，Load）：
- Extract：從硬碟上讀取資料 ------ 可以是本地（HDD 或 SSD），也可以是網盤（GCS 或 HDFS）
- Transform：使用 CPU 去解析、預處理資料 ------ 比如：影象解碼、資料增強、變換（比如：隨機裁剪、翻轉、顏色變換）、打亂、batching。
- Load：將 Transform 後的資料載入到 計算裝置 ------ 例如：GPU、TPU 等裝置。
上述的資料輸入管道使用 CPU 來進行資料的 ETL 過程，從而讓 GPU、TPU 等裝置專心進行模型的訓練過程（提高了裝置的利用率）。另外，將資料輸入管道抽象為 ETL 過程，有利於我們對資料輸入管道進行優化。
當使用 tf.estimator.Estimator API 時，input_fn 需要完成 Extract 和 Transform 兩個階段。
def parse_fn(example):
"Parse TFExample records and perform simple data augmentation."
example_fmt = {
"img_encoded": tf.FixedLenFeature((), tf.string, ""),
"img_label": tf.FixedLenFeature((), tf.int64, -1)
}
parsed = tf.parse_single_example(example, example_fmt)
image = tf.image.decode_image(parsed["img_encoded"])
return image, parsed["img_label"]
def input_fn(batch_size):
files = tf.data.Dataset.list_files("/path/to/dataset/train-.tfrecord")
ds = files.interleave(tf.data.TFRecordDataset, cycle_length=1)
ds = ds.shuffle(buffer_size=batch_size4)
ds = ds.map(map_func=parse_fn)
ds = ds.batch(batch_size=batch_size)
return ds
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
- 15
- 16
- 17
注：後面的過程以上面的 input_fn 為基礎。
2. 資料輸入管道的效能優化
新計算裝置使得網路的訓練越來越快，所以我們必須細心地設計 CPU 上執行的資料輸入管道（防止其成為系統的效能瓶頸）。tf.data 提供了資料輸入管道所需的各種部件，藉助其，我們可以實現高效的資料輸入管道，優化 ETL 過程的各個步驟。
2.1 資料準備(ET)和消耗過程(L)的解耦 ------ prefetch解耦、重疊兩過程
在執行一個訓練 step 之前，你必須 Extract、Transform 訓練資料，然後將它饋送給計算裝置。在以前，當 CPU 為計算準備資料時，計算裝置處於閒置狀態；當計算裝置執行訓練 step 時，CPU 處於閒置狀態。因此，單個訓練 step 的時間等於 CPU 準備資料的時間 + 計算裝置執行訓練 step 的時間。
Pipelining 將訓練 step 中的 資料準備 和 模型執行 “並行”。當計算裝置在執行第 N 個訓練 step 時，CPU 為第 N+1 個訓練 step 準備資料。通過兩個過程的重疊，單個訓練 step 的時間等於 CPU 準備資料的時間 和 計算裝置執行訓練 step 的時間中較大值。
tf.data.Dataset.prefetch 提供了 software pipelining 機制。該函式解耦了 資料產生的時間 和 資料消耗的時間。具體來說，該函式有一個後臺執行緒和一個內部快取區，在資料被請求前，就從 dataset 中預載入一些資料（進一步提高效能）。prefech(n) 一般作為最後一個 transformation，其中 n 為 batch_size。
prefetch 的使用方法如下：
dataset = dataset.batch(batch_size=FLAGS.batch_size)
dataset = dataset.prefetch(buffer_size=FLAGS.prefetch_buffer_size) # last transformation
return dataset
- 1
- 2
- 3
注意：只要你的 資料產生過程 和 資料消耗過程 可以重合（無論重合多少），那麼 prefetch 就能為你帶來效能提升。
2.2 資料變換(T)的並行化 ------ 並行map，融合map和batch
使用 tf.data.Dataset.map，我們可以很方便地對資料集中的各個元素進行預處理。因為輸入元素之間時獨立的，所以可以在多個 CPU 核心上並行地進行預處理。map 變換提供了一個 num_parallel_calls 引數去指定並行的級別。例如，下圖為 num_parallel_calls=2 時 map 變換的示意圖：
num_parallel_calls 引數的最優值取決於你的硬體、訓練資料的特質（比如：它的 size、shape）、map 函式的計算量 和 CPU 上同時進行的其它處理。比較簡單的一個設定方法是：將 num_parallel_calls 設定為 CPU 的核心數。例如，CPU 有四個核心時，將 num_parallel_calls 設定為 4 將會很高效。相反，如果 num_parallel_calls 大於 CPU 的核心數，將導致低效的排程，導致輸入管道的效能下降。
map 變換開啟並行化的方法如下：
dataset = dataset.map(map_func=parse_fn, num_parallel_calls=FLAGS.num_parallel_calls)
- 1
另外，如果你的 batch_size 比較大（成百上千），以 batch 的形式進行並行能夠帶來額外的效能提高。為此，tf.data 提供了 tf.contrib.data.map_and_batch 函式，其高效地融合了 map、batch 兩個變換。
為了融合 map 和 batch 兩個變換，我們只需要將：
dataset = dataset.map(map_func=parse_fn, num_parallel_calls=FLAGS.num_parallel_calls)
dataset = dataset.batch(batch_size=FLAGS.batch_size)
- 1
- 2
改為：
dataset = dataset.apply(tf.contrib.data.map_and_batch(
map_func=parse_fn, batch_size=FLAGS.batch_size))
- 1
- 2
2.3 資料讀取(E)的並行化------並行地讀取並解析多個數據檔案
在實際應用中，輸入資料可能被儲存在網盤（例如，GCS 或 HDFS）（要麼因為輸入資料不適合本地，要麼因為訓練是分散式的，在每臺機器上覆制輸入資料是沒有意義的）。另外，在本地能夠很好的讀取資料的資料輸入管道也可能會卡在 I/O 瓶頸上，因為 本地 和 遠端儲存 有以下區別：
- Time-to-first-byte（讀取第一個bytes的時間）：從遠端儲存讀取檔案的第一個位元組的時間比本地儲存長一個數量級。
- Read throughput（讀取吞吐量）：雖然遠端儲存通常提供大的聚合頻寬，但是讀取單個檔案可能僅能利用該頻寬的一小部分。
另外，一旦原始位元組被讀取到記憶體中，也可能需要對資料進行反序列化或解密（例如：protobuf），這將導致額外的負載。不管資料是本地儲存還是遠端儲存，該開銷都存在，但如果資料未被高效地預載入，則遠端情況下可能更糟。
為了減輕各種資料讀取(E)開銷的影響，tf.data 提供了 tf.contrib.data.parallel_interleave 函式。該函式可以並行地從多個檔案中提取並解析資料。同時讀取的檔案數可以通過引數 cycle_length 來指定。
下圖說明了將 parallel_interleave 中的 cycle_length=2 時的效果：
為了並行地讀取資料(E)，只需將：
dataset = files.interleave(tf.data.TFRecordDataset) # 現在該函式已經加了cycle_length引數
- 1
改為：
dataset = files.apply(tf.contrib.data.parallel_interleave(
tf.data.TFRecordDataset, cycle_length=FLAGS.num_parallel_readers))
- 1
- 2
遠端儲存系統的吞吐量會受負載、網路事件等影響。為了緩解這種影響，可以將 parallel_interleave 和 prefetching 結合使用（詳情見：tf.contrib.data.parallel_interleave）
預設情況下，parallel_interleave 函式為元素提供一個確定性順序，以方便再現。作為 prefetching 的一個替代方案（這在某些情況下，可能不高效），parallel_interleave 變換也提供了一個選項去提高效能（代價是元素的順序的確定性）。尤其是，如果 sloppy 引數被設定為 True，變換可能偏離設定的順序，通過臨時跳過在下一個元素被請求時元素不可用的檔案。
3. 資料輸入管道效能的進一步優化
tf.data API 是圍繞可組合的變換設計的（為使用者提供靈活性）。雖然這些變換中的很多變換的次序是可交換的，但某些變換的次序對效能有影響。
3.1 Map and Batch ------ map開銷很小時，batch形式的map更高效
將使用者自定義的函式傳給 map 函式 會產生排程、執行使用者自定義函式的負載。一般情況下，這個負載與自定義函式的計算量相比很小。但是，如果 map 的函式的計算量很小，這個負載將是主要開銷。在這種情況下，我們推薦使用向量化的自定義函式（它一次對一個batch進行變換），並且在 map 變換前使用 batch 變換。
3.2 Map and Cache ------ 通過cache進一步加速
tf.data.Dataset.cache 變化能夠在記憶體或本地儲存器上快取一個數據集。如果傳遞給 map 變換的使用者自定義函式的計算量很大，只要得到的資料集仍然適合記憶體或本地儲存，就可以在 map 轉換之後應用 cache 轉換。
如果使用者定義函式導致儲存資料集需要的空間超過了 cache 的容量，考慮提前對資料集進行預處理，以減少資源的使用。
注意：cache將資料集進行快取能夠有效地提高資料輸入管道的效能，但是，cache位置放置錯誤時，會導致模型效能下降。
3.3 Map and Interleave / Prefetch / Shuffle ------ 變換的順序對記憶體使用量的影響
因為各個變換函式（包括 interleave，prefetch，shuffle）都有自己的內部快取，所以如果傳給 map 變換的 使用者自定義函式 改變了元素的 size，那麼 map 變換的次序影響記憶體的使用量。通常情況下，我們建議選擇記憶體使用量更低的次序，除非不同的次序能夠產生效能上的提高（例如，為了使用融合的 tf.contrib.data.map_and_batch）。
3.4 Repeat and Shuffle ------ repeat在前效能優，shuffle在前次序強
tf.data.Dataset.repeat 變換重複輸入資料有限次（或無限次）；資料的每一次重複稱為一個 epoch。tf.data.Dataset.shuffle 變換隨機打亂資料集 example 的次序。
如果 repeat 變換被放在 shuffle 變換之前，那麼 epoch 邊界將變得模糊。也就是說，某些元素可以在其他元素出現一次之前重複。另一方面，如果在 repeat 變換之前應用 shuffle 變換，那麼在每個 epoch 開始時，效能可能會下降（因為這時，也需要進行 shuffle 變化的初始化）。換句話說，將 repeat 放置在 shuffle 之前，提供了更好的效能，將 shuffle 放置在 repeat 之前，提供了更強的次序保證。
當可能時，我們推薦使用融合op：tf.contrib.data.shuffle_and_repeat 變換，這個變換在效能和更強的次序保證上都是最好的（good performance and strong ordering guarantees）。否則，我們推薦在 repeat 之前使用 shuffle。
4. 資料輸入管道的最優實現
下面是設計最優資料輸入管道的建議：
- 使用
prefetch函式去重疊 資料讀取器 和 資料消耗器的工作。我們尤其推薦在輸入管道的末端新增prefetch(n)（n是batch size），以重疊 CPU 上的變換 及 GPU/TPU裝置上的訓練。詳見【2.1】 - 通過設定
num_parallel_calls引數，來並行 map 變換。我們建議使用將該引數設定為 CPU 的核心數。詳見【2.2】 - 如果你使用 batch 變換來將預處理好的元素 batching，我們建議使用融合op：
map_and_batch變換；尤其是你如果使用大的batch size。詳見【2.2】 - 如果你的資料存在遠端儲存上，（且有時需要解析），我們建議使用
parallel_interleave來並行資料的讀取和解析。詳見【2.3】 - 將簡單的使用者自定義函式進行向量化，然後傳遞給 map 變換去分攤 使用者自定義函式有關的呼叫、執行的負載。詳見【3.1】
- 如果你的資料能夠載入到記憶體，使用
cache變化去在訓練的第一個 epoch 將資料集快取到記憶體，所以能避免後來的 epoch 讀取、解析、變換資料的負載。詳見【3.2】 - 如果你的預處理會增加你資料的 size，我們建議你首先使用
interleave、prefetch、shuffle變換去減少記憶體使用量（如果可能）。詳見【3.3】 - 我們建議在
repeat變換之前使用shuffle變換，最好使用融合op：shuffle_and_repeat變換。詳見【3.4】
英文版本見：https://tensorflow.google.cn/performance/datasets_performance
</div>
<link href="https://csdnimg.cn/release/phoenix/mdeditor/markdown_views-a47e74522c.css" rel="stylesheet">
</div>
<script>
(function(){
function setArticleH(btnReadmore,posi){
var winH = $(window).height();
var articleBox = $("div.article_content");
var artH = articleBox.height();
if(artH > winH*posi){
articleBox.css({
'height':winH*posi+'px',
'overflow':'hidden'
})
btnReadmore.click(function(){
if(typeof window.localStorage === "object" && typeof window.csdn.anonymousUserLimit === "object"){
if(!window.csdn.anonymousUserLimit.judgment()){
window.csdn.anonymousUserLimit.Jumplogin();
return false;
}else if(!currentUserName){
window.csdn.anonymousUserLimit.updata();
}
}
articleBox.removeAttr("style");
$(this).parent().remove();
})
}else{
btnReadmore.parent().remove();
}
}
var btnReadmore = $("#btn-readmore");
if(btnReadmore.length>0){
if(currentUserName){
setArticleH(btnReadmore,3);
}else{
setArticleH(btnReadmore,1.2);
}
}
})()
</script>
TensorFlow版本：1.12.0