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

時間 2019-11-17

標籤 tf.data data 優化訓練數據輸入管道 google 開發者大會欄目 Google 简体版

原文原文鏈接

Google 開發者大會 (Google Developer Days，簡稱 GDD) 是展現 Google 最新開發者產品和平臺的全球盛會，旨在幫助你快速開發優質應用，發展和留住活躍用戶羣，充分利用各類工具得到更多收益。2018 Google 開發者大會於 9 月 20 日和 21 日於上海舉辦。👉Google 開發者大會 2018 掘金專題python

GDD 2018 次日的 9 月 21 日，陳爽（Google Brain 軟件工程師）爲咱們帶來了《以 tf.data 優化訓練數據》，講解如何使用 tf.data 爲各種模型打造高性能的 TensorFlow 輸入渠道，本文將摘錄演講技術乾貨。git

數據輸入管道

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

ETL 系統

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

tf.data：爲機器學習設計的數據輸入系統

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

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 性能指南：sql

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 提供現有高效配方

上圖能夠簡單替換爲一個函數：app

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)
複製代碼