Google 開發者大會 (Google Developer Days，簡稱 GDD) 是展示 Google 最新開發者產品和平臺的全球盛會，旨在幫助你快速開發優質應用，發展和留住活躍使用者群，充分利用各種工具獲得更多收益。2018 Google 開發者大會於 9 月 20 日和 21 日於上海舉辦。?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 這個兼具高效丶靈活與易用的 API，同時瞭解如何運用管道化及其他優化手段來增進運算效能，以及許多可能未曾發現的實用函式。
資源
- 入門指南
- 示例程式碼