以 tf.data 優化訓練資料輸入管道丨 Google 開發者大會 2018

谷歌 · 發表 2018-09-21 20:47:56

摘要： Google 開發者大會 (Google Developer Days，簡稱 GDD) 是展示 Google 最新開發者產品和平臺的全球盛會，旨在幫助你快速開發優質應用，發展和留住活躍使用者群，充分利用各種工具獲得更多收益。2018 Google 開發者大會於 9 月 20 日和 21 日於...

Google 開發者大會 (Google Developer Days，簡稱 GDD) 是展示 Google 最新開發者產品和平臺的全球盛會，旨在幫助你快速開發優質應用，發展和留住活躍使用者群，充分利用各種工具獲得更多收益。2018 Google 開發者大會於 9 月 20 日和 21 日於上海舉辦。:point_right: ofollow,noindex">Google 開發者大會 2018 掘金專題

GDD 2018 第二天的 9 月 21 日，陳爽（Google Brain 軟體工程師）為我們帶來了《以 tf.data 優化訓練資料》，講解如何使用 tf.data 為各類模型打造高效能的 TensorFlow 輸入渠道，本文將摘錄演講技術乾貨。

資料輸入管道

大多人將時間和金錢花在神經網路架構上，資料輸入容易被忽略
沒有好的資料輸入管道，GPU 再強速度也不會顯著提高
目標：高效丶靈活丶易用

ETL 系統

提取資料（Extract）：將訓練資料從存取器（硬碟丶雲端等）提取
轉換資料（Transform）：將資料轉換為模型可讀取的資料，同時進行資料清洗等預處理
裝載資料（Load）：將處理好的資料裝載至加速器

tf.data：為機器學習設計的資料輸入系統

圖中程式碼分別對應 ETL 系統的三個步驟，使用 tf.data 即可輕鬆實現。

tf.data 優化手段：以上圖程式碼為例

多執行緒處理（使用 num_parallel_reads ）

files = tf.data.Dataset.list_files("training-*-of-1024.tfrecord")
dataset = tf.data.TFRecordDataset(files, num_parallel_reads=32)
複製程式碼

合併轉換步驟（如 shuffle_and_repaeat , map_and_batch ）

dataset = dataset.apply(tf.contrib.data.shuffle_and_repaeat(10000, NUM_EPOCHS))
dataset = dataset.apply(tf.contrib.data.map_and_batch(lambda x: ..., BATCH_SIZE))
複製程式碼

流水線化（使用 prefetch_to_device ）

dataset = dataset.apply(tf.contrib.data.prefetch_to_device("/gpu:0"))
複製程式碼

最終程式碼如下圖所示，更多優化手段可以參考tf.data 效能指南：

tf.data 的靈活性

支援函數語言程式設計

如上圖，可以用自定義的 map_fn 處理 TensorFlow 或相容的函式，同時支援 AutoGraph 處理過的函式。

支援不同語言與資料型別

使用 Dataset.form_generator() 支援 Python 程式碼生成 Dataset
使用 DatasetOpKernel 和 tf.load_op_library 支援自定義 C++ 資料處理程式碼

如下圖，使用 Python 自帶的 urllib 獲取伺服器資料，存入 dataset：

支援多種資料來源

如普通檔案系統丶GCP 雲儲存丶其他雲儲存丶SQL 資料庫等。

讀取 Google 雲儲存的 TFRecord 檔案示例：

files = tf.contrib.data.TFRecordDataset(
"gs://path/to/file.tfrecord", num_parallel_reads=32)
複製程式碼

使用自訂 SQL 資料庫示例：

files = tf.contrib.data.SqlDataset(
"sqllite", "/foo/db.sqlite", "SELECT name, age FROM people", 
(tf.string, tf.int32))
複製程式碼

tf.data 的易用性

在 Eager 執行模式下，可以直接使用 Python for 迴圈：

tf.enable_eager_execution()
for batch in dataset:
train_model(batch)
複製程式碼

為 TF Example 或 CSV 提供現有高效配方

可以簡單替換為一個函式：

dataset = tf.contrib.data.make_batched_features_dataset(
"training-*-of-1024.tfrecord",
BATCH_SIZE, features, num_epochs=NUM_EPOCHS)
複製程式碼

使用 CSV 資料集的情境：

dataset = tf.contrib.data.make_csv_dataset(
"*.csv", BATCH_SIZE, num_epochs=NUM_EPOCHS)
複製程式碼

使用 AUTOTUNE 自動調節管道

可以簡單的使用 AUTOTUNE 找到 prefetching 的最佳引數：

dataset = dataset.prefetch(tf.contrib.data.AUTOTUNE)
複製程式碼

支援 Keras 和 Estimators 相互相容

對於 Keras，可以將 dataset 直接傳遞使用；對於 Estimators 訓練函式，將 dataset 包裝至輸入函式並返回即可，如下示例：

def input_fn():
dataset = tf.contrib.data.make_csv_dataset(
"*.csv", BATCH_SIZE, num_epochs=NUM_EPOCHS)
return dataset

tf.estimator.Estimator(model_fn=train_model).train(input_fn=input_fn)
複製程式碼

實際運用經驗

原始 tf.data 資料輸入程式碼： ~150 影象 / 秒
管道化的 tf.data 資料輸入程式碼： ~1,750 影象 / 秒 => 12倍的效能！
Cloud TPU 上使用 tf.data： ~4,100 影象 / 秒
Cloud TPU Pod 上使用 tf.data： ~219,000 影象 / 秒

以 tf.data 優化訓練資料輸入管道 丨 Google 開發者大會 2018