TPU 速度資料管道:tf.data.Dataset 和 TFRecords

1. 總覽

TPU 速度飛快,訓練資料串流必須跟上訓練速度。在本實驗室中,您將瞭解如何使用 tf.data.Dataset API 從 GCS 載入資料,以便提供給 TPU。

本實驗室是「Keras 在 TPU 上運作」系列的 Part 1。您可以依照下列順序執行這些操作,也可以獨立執行。

ca8cc21f6838eccc.png

課程內容

  • 如何使用 tf.data.Dataset API 載入訓練資料
  • 使用 TFRecord 格式,以便從 GCS 有效率地載入訓練資料

意見回饋

如果您發現這個程式碼研究室有任何不妥之處,請告訴我們。您可以透過 GitHub 問題提供意見 [feedback link]。

2. Google Colaboratory 快速入門

這個實驗室使用 Google Collaboratory,您不需要進行任何設定。Colaboratory 是用於教育用途的線上筆記平台。提供免費 CPU、GPU 和 TPU 訓練。

688858c21e3beff2.png

您可以開啟這個範例筆記本並執行多個儲存格,熟悉 Colaboratory。

c3df49e90e5a654f.png Welcome to Colab.ipynb

選取 TPU 後端

8832c6208c99687d.png

在 Colab 選單中,依序選取「執行階段」>「變更執行階段類型」,然後選取 TPU。在本程式碼研究室中,您將使用強大的 TPU (Tensor Processing Unit),以便進行硬體加速訓練。系統會在第一次執行時自動連線至執行階段,您也可以使用右上角的「連線」按鈕。

執行筆記本

76d05caa8b4db6da.png

按一下儲存格,然後使用 Shift + Enter 鍵,一次執行一個儲存格。您也可以透過「執行階段」>「執行全部」執行整個筆記本

Table of contents

429f106990037ec4.png

所有筆記本都有目錄。您可以使用左側的黑色箭頭開啟。

隱藏的儲存格

edc3dba45d26f12a.png

部分儲存格只會顯示標題。這是 Colab 專屬的筆記本功能。您可以按兩下這些程式碼,以查看裡面的程式碼,但不太有趣。通常支援或視覺化函式。您仍需要執行這些儲存格,才能定義其中的函式。

驗證

cdd4b41413100543.png

只要您使用授權帳戶進行驗證,Colab 就能存取您的私人 Google Cloud Storage 值區。上方的程式碼片段會觸發驗證程序。

3. [INFO] 什麼是 Tensor Processing Unit (TPU)?

概述

f88cf6facfc70166.png

這段程式碼會在 Keras 的 TPU 上訓練模型 (如果無法使用 TPU,則會改回使用 GPU 或 CPU):

try: # detect TPUs
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect()
    strategy = tf.distribute.TPUStrategy(tpu)
except ValueError: # detect GPUs
    strategy = tf.distribute.MirroredStrategy() # for CPU/GPU or multi-GPU machines

# use TPUStrategy scope to define model
with strategy.scope():
  model = tf.keras.Sequential( ... )
  model.compile( ... )

# train model normally on a tf.data.Dataset
model.fit(training_dataset, epochs=EPOCHS, steps_per_epoch=...)

今天我們會使用 TPU,以互動速度 (每次訓練執行作業的分鐘數) 建構及最佳化花朵分類器。

688858c21e3beff2.png

為什麼要使用 TPU?

現代 GPU 以可程式設計的「核心」為基礎,這種非常靈活的架構可處理各種工作,例如 3D 算繪、深度學習、物理模擬等。另一方面,TPU 則是將傳統向量處理器與一個專用矩陣相乘單元配對,在所有大型矩陣乘積佔據的主要任務 (例如類神經網路) 的任務上都表現出優異。

8eb3e718b8e2ed08.png

插圖:密集的類神經網路層以矩陣乘法運算為矩陣,會一次透過類神經網路處理一批八張圖像,請透過一行 x 資料欄的乘法驗證,確認這確實是圖片中所有像素值的加權總和。卷積層也可以以矩陣相乘表示,但會比較複雜 (這裡有說明,請參閱第 1 節)。

硬體

MXU 和 VPU

TPU v2 核心是由矩陣乘積單元 (MXU) 組成,後者會執行矩陣乘法和向量處理器 (VPU),可執行啟動、softmax 等所有其他工作。VPU 會處理 float32 和 int32 運算。另一方面,MXU 會以混合精度 16-32 位元浮點格式運作。

7d68944718f76b18.png

混合精確度浮點和 bfloat16

MXU 使用 bfloat16 輸入與 float32 輸出來計算矩陣乘積。中階累計作業是以 float32 精度計算。

19c5fc432840c714.png

類神經網路訓練通常足以應付低浮點精確度所產生的雜訊。在某些情況下,雜訊甚至能協助最佳化工具收縮。通常使用 16 位元浮點精確度來加速運算,但 float16 和 float32 格式有非常不同的範圍。將精確度從 float32 降為 float16,通常會產生過流和反向溢位。現有解決方案存在,但通常需要額外工作才能讓 float16 運作。

因此,Google 在 TPU 中推出了 bfloat16 格式。bfloat16 是經過截斷的 float32,其指數位元和範圍與 float32 完全相同。這超出了 TPU 計算矩陣乘法的精度和 bfloat16 輸入內容,但使用 float32 輸出,通常都不需要變更程式碼,就能享有低精確度的效能提升。

Systolic array

MXU 使用所謂的「脈動陣列」架構,在硬體中實作矩陣乘積,資料元素會在硬體運算單元陣列中流動。(在醫學中,「脈衝」是指心臟收縮和血流,此處指向資料的流動)。

矩陣乘法的基本元素是介於一個矩陣的線條與另一個矩陣的欄之間的點積 (請參閱本節上方的插圖)。對於矩陣相乘 Y=X*W,結果的一個元素會是:

Y[2,0] = X[2,0]*W[0,0] + X[2,1]*W[1,0] + X[2,2]*W[2,0] + ... + X[2,n]*W[n,0]

在 GPU 上,其中一個會將這個內積分裝成 GPU 的「核心」,然後在盡可能平行執行大量「核心」上執行該積點,嘗試同時計算產生的矩陣的每個值。如果產生的矩陣為 128x128 大,則需要 128x128=16K 的「核心」才可使用,而這通常不可能。最大的 GPU 大約有 4,000 個核心。另一方面,TPU 在 MXU 中運算單元使用的硬體最低只有 bfloat16 x bfloat16 => float32 乘數加總,除此之外不用其他。這些都太小,TPU 可以在 128x128 MXU 中實作 16K 的容器,並一次處理這個矩陣乘法。

f1b283fc45966717.gif

插圖:MXU 收縮壓陣列,運算元素為乘積。一個矩陣的值會載入陣列 (紅點)。其他矩陣的值流經陣列 (灰點)。垂直線會將值向上傳播。水平線會傳播部分總和。這之後是使用者的練習,用來驗證資料通過陣列時,你得到的矩陣乘積結果來自右側。

此外,雖然內積是以 MXU 計算中度積,但在相鄰的運算單元之間,中繼總和其實只會流動。不需要儲存和擷取記憶體,甚至是註冊檔案。最終結果是 TPU 脈動陣列架構在計算矩陣乘法時,具有顯著密度和效能優勢;此外,在計算矩陣乘法時,TPU 的速度優勢也不可忽略。

Cloud TPU

在 Google Cloud Platform 上要求一個「Cloud TPU v2」時,您會取得一個含有 PCI 連接 TPU 板的虛擬機器 (VM)。TPU 電路板有四個雙核心 TPU 晶片。每個 TPU 核心都配備 VPU (向量處理單元) 和 128x128 MXU (矩陣乘法單元)。這類「Cloud TPU」通常會透過網路連線至要求的 VM。完整的圖片如下所示:

dfce5522ed644ece.png

插圖:您的 VM 具有附加網路的「Cloud TPU」加速器,「Cloud TPU」本身就是 VM 配上 PCI 連接的 TPU 電路板,上面有四個雙核心 TPU 晶片。

TPU Pod

在 Google 的資料中心,TPU 會連線至高效能運算 (HPC) 互連網路,因此看起來就像一個極大的加速器。Google 會呼叫 Pod,且最多可包含 512 個 TPU v2 核心或 2048 個 TPU v3 核心。

2ec1e0d341e7fc34.jpeg

插圖:TPU v3 Pod。透過 HPC 互連網路連線的 TPU 主機板和機架。

在訓練期間,系統會使用全減算演算法在 TPU 核心之間交換梯度 (這裡有全減算的詳細說明)。訓練中的模型可透過大量批次訓練來充分利用硬體。

d97b9cc5d40fdb1d.gif

圖解:在 Google TPU 的 2-D 環面網狀 HPC 網路使用 all-red 演算法,在訓練期間同步處理梯度。

軟體

大批次訓練

TPU 的理想批次大小為每個 TPU 核心 128 個資料項目,但硬體已能顯示每個 TPU 核心 8 個資料項目有良好的使用率。請記住,一個 Cloud TPU 有 8 個核心。

在本程式碼研究室中,我們將使用 Keras API。在 Keras 中,您指定的批次是整個 TPU 的全球批次大小。批次將自動分割為 8 個,並在 TPU 的 8 核心上執行。

da534407825f01e3.png

如需其他效能提示,請參閱 TPU 效能指南。如果批次大小非常大,某些模型可能需要特別處理,詳情請參閱 LARSOptimizer

後製:XLA

TensorFlow 程式會定義運算圖形。TPU 不會直接執行 Python 程式碼,而是執行 Tensorflow 程式定義的運算圖形。理論上,名為 XLA 的編譯器 (加速線性代數編譯器) 會將運算節點的 Tensorflow 圖形轉換為 TPU 機器程式碼。這個編譯器也會對程式碼和記憶體配置執行許多進階最佳化。系統會在工作傳送至 TPU 時自動編譯。您不必在建構鏈結中明確加入 XLA。

edce61112cd57972.png

範例:為了在 TPU 上執行,TensorFlow 程式定義的運算圖形會先轉譯為 XLA (加速線性代數編譯器) 表示法,然後由 XLA 編譯為 TPU 機器碼。

在 Keras 中使用 TPU

自 TensorFlow 2.1 起,Keras API 就支援 TPU。Keras 支援 TPU 和 TPU Pod。以下是適用於 TPU、GPU 和 CPU 的範例:

try: # detect TPUs
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect()
    strategy = tf.distribute.TPUStrategy(tpu)
except ValueError: # detect GPUs
    strategy = tf.distribute.MirroredStrategy() # for CPU/GPU or multi-GPU machines

# use TPUStrategy scope to define model
with strategy.scope():
  model = tf.keras.Sequential( ... )
  model.compile( ... )

# train model normally on a tf.data.Dataset
model.fit(training_dataset, epochs=EPOCHS, steps_per_epoch=...)

在這個程式碼片段中:

  • TPUClusterResolver().connect() 會在網路中找到 TPU。在大多數 Google Cloud 系統 (AI 平台工作、Colaboratory、Kubeflow、透過「ctpu up」公用程式建立的深度學習 VM) 上,無需參數即可運作。這些系統會透過 TPU_NAME 環境變數得知 TPU 的位置。如要手動建立 TPU,請在目前使用的 VM 上設定 TPU_NAME env. var,或是使用明確參數呼叫 TPUClusterResolverTPUClusterResolver(tp_uname, zone, project)
  • TPUStrategy 是實作分發和「all-reduce」梯度同步處理演算法的部分。
  • 系統是透過範圍套用策略。模型必須在策略 scope() 內定義。
  • tpu_model.fit 函式預期輸入用於 TPU 訓練的 tf.data.Dataset 物件。

常見的 TPU 移植工作

  • 在 TensorFlow 模型中載入資料的方法很多,但如果是 TPU,則必須使用 tf.data.Dataset API。
  • TPU 速度飛快,且執行時擷取資料經常成為瓶頸。請參閱 TPU 效能指南,瞭解您可以使用哪些工具偵測資料瓶頸和其他效能提示。
  • 系統會將 int8 或 int16 號碼視為 int32。TPU 沒有在 32 位元以下運作的整數硬體。
  • 不支援部分 Tensorflow 作業。請參閱這份清單。好消息是,這項限制只適用於訓練程式碼,也就是透過模型的正向和反向傳遞。您仍可在資料輸入管道中使用所有 TensorFlow 作業,因為 TensorFlow 將在 CPU 上執行。
  • TPU 不支援 tf.py_func

4. 載入資料

c0ecb860e4cad0a9.jpeg cc4781a7739c49ae.jpeg 81236b00f8bbf39e.jpeg 961e2228974076bb.jpeg 7517dc163bdffcd5.jpeg 96392df4767f566d.png

我們會使用花朵圖片資料集。目標是學習將這些花朵分類為 5 種花型。資料載入作業是使用 tf.data.Dataset API 執行。首先,讓我們瞭解這個 API。

實作

請開啟下列筆記本,執行儲存格 (Shift + ENTER),並按照「WORK REQUIRED」標籤中的指示操作。

c3df49e90e5a654f.png Fun with tf.data.Dataset (playground).ipynb

其他資訊

關於「flowers」資料集

資料集儲存在 5 個資料夾中。每個資料夾都包含一種花朵。資料夾名稱分別為「sunflowers」、「daisy」、「dandelion」、「tulips」和「roses」。這些資料託管於 Google Cloud Storage 的公開值區。摘錄:

gs://flowers-public/sunflowers/5139971615_434ff8ed8b_n.jpg
gs://flowers-public/daisy/8094774544_35465c1c64.jpg
gs://flowers-public/sunflowers/9309473873_9d62b9082e.jpg
gs://flowers-public/dandelion/19551343954_83bb52f310_m.jpg
gs://flowers-public/dandelion/14199664556_188b37e51e.jpg
gs://flowers-public/tulips/4290566894_c7f061583d_m.jpg
gs://flowers-public/roses/3065719996_c16ecd5551.jpg
gs://flowers-public/dandelion/8168031302_6e36f39d87.jpg
gs://flowers-public/sunflowers/9564240106_0577e919da_n.jpg
gs://flowers-public/daisy/14167543177_cd36b54ac6_n.jpg

為什麼要使用 tf.data.Dataset?

Keras 和 TensorFlow 的所有訓練和評估函式都會接受 Dataset。在資料集中載入資料後,API 會提供所有常見功能,這些功能可用於訓練神經網路的資料:

dataset = ... # load something (see below)
dataset = dataset.shuffle(1000) # shuffle the dataset with a buffer of 1000
dataset = dataset.cache() # cache the dataset in RAM or on disk
dataset = dataset.repeat() # repeat the dataset indefinitely
dataset = dataset.batch(128) # batch data elements together in batches of 128
AUTOTUNE = tf.data.AUTOTUNE
dataset = dataset.prefetch(AUTOTUNE) # prefetch next batch(es) while training

如要瞭解效能提示和資料集最佳做法,請參閱這篇文章。參考說明文件請見這裡

tf.data.Dataset 基本概念

資料通常會以多個檔案,在這裡的圖片。您可以透過以下方式建立檔案名稱資料集:

filenames_dataset = tf.data.Dataset.list_files('gs://flowers-public/*/*.jpg')
# The parameter is a "glob" pattern that supports the * and ? wildcards.

接著,您可以為每個檔案名稱「對應」函式,通常會將檔案載入記憶體並解碼為實際資料:

def decode_jpeg(filename):
  bits = tf.io.read_file(filename)
  image = tf.io.decode_jpeg(bits)
  return image

image_dataset = filenames_dataset.map(decode_jpeg)
# this is now a dataset of decoded images (uint8 RGB format)

如要對資料集進行疊代作業:

for data in my_dataset:
  print(data)

元組資料集

在監督式學習中,訓練資料集通常是以成對的訓練資料與正確答案組成。為此,解碼函式可能會傳回元組。如此一來,當您對其進行疊代時,就會傳回包含元組和元組的資料集。回傳的值是 Tensorflow 張量,可供模型使用。您可以對這些值呼叫 .numpy(),查看原始值:

def decode_jpeg_and_label(filename):
  bits = tf.read_file(filename)
  image = tf.io.decode_jpeg(bits)
  label = ... # extract flower name from folder name
  return image, label

image_dataset = filenames_dataset.map(decode_jpeg_and_label)
# this is now a dataset of (image, label) pairs 

for image, label in dataset:
  print(image.numpy().shape, label.numpy())

結論:逐一載入圖片的速度很慢!

疊代這個資料集時,您會發現每秒可以載入 1 到 2 張圖片。太慢了!我們將用於訓練的硬體加速器可維持數倍的速度。請參閱下一節,瞭解我們如何達成此目標。

解決方案

以下是解決方案的筆記本。如果遇到問題,可以使用此功能。

c3df49e90e5a654f.png Fun with tf.data.Dataset (solution).ipynb

涵蓋內容

  • 🤔? tf.data.Dataset.list_files
  • 🤔 tf.data.Dataset.map
  • 🤔 元組的資料集
  • 😀 逐一處理資料集

請花點時間研讀這份檢查清單,

5. 快速載入資料

本研究室使用的 Tensor Processing Unit (TPU) 硬體加速器非常快速。資料提供的速度要夠快,才能保持繁忙度,是該公司的一大挑戰。Google Cloud Storage (GCS) 可維持極高的傳輸量,但與所有雲端儲存系統一樣,啟動連線時需要進行一些網路來回傳輸。因此,將資料儲存為數千個個別檔案並不理想。我們將在較少的檔案中批次處理這些檔案,並使用 tf.data.Dataset 功能,平行讀取多個檔案。

詳閱

以下筆記本會載入圖片檔、將其調整為常見大小,然後儲存至 16 個 TFRecord 檔案,程式碼位於下列筆記本中。請快速閱讀。您不必執行這項操作,因為程式碼研究室的其他部分會提供正確的 TFRecord 格式資料。

c3df49e90e5a654f.png Flower pictures to TFRecords.ipynb

理想的資料配置,可獲得最佳 GCS 吞吐量

TFRecord 檔案格式

Tensorflow 偏好的資料檔案格式是以 protobuf 為基礎的 TFRecord 格式。其他序列化格式也適用,但您可以直接透過以下方式從 TFRecord 檔案載入資料集:

filenames = tf.io.gfile.glob(FILENAME_PATTERN)
dataset = tf.data.TFRecordDataset(filenames)
dataset = dataset.map(...) # do the TFRecord decoding here - see below

為達到最佳效能,建議您使用以下較複雜的程式碼,一次讀取多個 TFRecord 檔案。此程式碼會並行讀取 N 個檔案,並忽略資料順序,以便加快讀取速度。

AUTOTUNE = tf.data.AUTOTUNE
ignore_order = tf.data.Options()
ignore_order.experimental_deterministic = False

filenames = tf.io.gfile.glob(FILENAME_PATTERN)
dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTOTUNE)
dataset = dataset.with_options(ignore_order)
dataset = dataset.map(...) # do the TFRecord decoding here - see below

TFRecord 一覽表

可將三種類型的資料儲存在 TFRecords 中:位元組字串 (位元組清單)、64 位元整數和 32 位元浮點。這些資料一律會儲存為清單,單一資料元素清單的大小為 1。您可以使用下列輔助函式,將資料儲存到 TFRecord。

寫入位元組字串

# warning, the input is a list of byte strings, which are themselves lists of bytes
def _bytestring_feature(list_of_bytestrings):
  return tf.train.Feature(bytes_list=tf.train.BytesList(value=list_of_bytestrings))

寫入整數

def _int_feature(list_of_ints): # int64
  return tf.train.Feature(int64_list=tf.train.Int64List(value=list_of_ints))

寫入浮點數

def _float_feature(list_of_floats): # float32
  return tf.train.Feature(float_list=tf.train.FloatList(value=list_of_floats))

編寫 TFRecord,使用上述輔助程式

# input data in my_img_bytes, my_class, my_height, my_width, my_floats
with tf.python_io.TFRecordWriter(filename) as out_file:
  feature = {
    "image": _bytestring_feature([my_img_bytes]), # one image in the list
    "class": _int_feature([my_class]),            # one class in the list
    "size": _int_feature([my_height, my_width]),  # fixed length (2) list of ints
    "float_data": _float_feature(my_floats)       # variable length  list of floats
  }
  tf_record = tf.train.Example(features=tf.train.Features(feature=feature))
  out_file.write(tf_record.SerializeToString())

如要從 TFRecord 讀取資料,您必須先宣告已儲存記錄的版面配置。在宣告中,您可以透過固定長度清單或變數長度清單存取任何已命名欄位:

從 TFRecords 讀取

def read_tfrecord(data):
  features = {
    # tf.string = byte string (not text string)
    "image": tf.io.FixedLenFeature([], tf.string), # shape [] means scalar, here, a single byte string
    "class": tf.io.FixedLenFeature([], tf.int64),  # shape [] means scalar, i.e. a single item
    "size": tf.io.FixedLenFeature([2], tf.int64),  # two integers
    "float_data": tf.io.VarLenFeature(tf.float32)  # a variable number of floats
  }

  # decode the TFRecord
  tf_record = tf.io.parse_single_example(data, features)

  # FixedLenFeature fields are now ready to use
  sz = tf_record['size']

  # Typical code for decoding compressed images
  image = tf.io.decode_jpeg(tf_record['image'], channels=3)

  # VarLenFeature fields require additional sparse.to_dense decoding
  float_data = tf.sparse.to_dense(tf_record['float_data'])

  return image, sz, float_data

# decoding a tf.data.TFRecordDataset
dataset = dataset.map(read_tfrecord)
# now a dataset of triplets (image, sz, float_data)

實用的程式碼片段:

讀取單一資料元素

tf.io.FixedLenFeature([], tf.string)   # for one byte string
tf.io.FixedLenFeature([], tf.int64)    # for one int
tf.io.FixedLenFeature([], tf.float32)  # for one float

讀取元素的固定大小清單

tf.io.FixedLenFeature([N], tf.string)   # list of N byte strings
tf.io.FixedLenFeature([N], tf.int64)    # list of N ints
tf.io.FixedLenFeature([N], tf.float32)  # list of N floats

讀取不同數量的資料項目

tf.io.VarLenFeature(tf.string)   # list of byte strings
tf.io.VarLenFeature(tf.int64)    # list of ints
tf.io.VarLenFeature(tf.float32)  # list of floats

VarLenFeature 會傳回稀疏向量,解碼 TFRecord 後需要執行其他步驟:

dense_data = tf.sparse.to_dense(tf_record['my_var_len_feature'])

您也可以在 TFRecords 中提供選用欄位。如果您在讀取欄位時指定預設值,那麼如果缺少該欄位,則會傳回預設值,而不是錯誤。

tf.io.FixedLenFeature([], tf.int64, default_value=0) # this field is optional

涵蓋內容

  • 🤔? 分割資料檔案以便快速存取 GCS
  • 😓 如何寫入 TFRecords。(您忘記語法了嗎?沒關係,將這個頁面加入書籤做為快速參考)
  • 🤔? 使用 TFRecordDataset 從 TFRecord 載入資料集

請花點時間在腦中過一遍這份檢查清單。

6. 恭喜!

您現在可以為 TPU 提供資料。請繼續進行下一個實驗室

實務應用 TPU

Cloud AI Platform 提供 TPU 和 GPU:

最後,我們非常喜歡使用者的意見。如果您發現本研究室有任何不妥之處,或認為有改進空間,請告訴我們。您可以透過 GitHub 問題提供意見 [feedback link]。

HR.png

Martin Görner ID small.jpg
作者:Martin Görner
Twitter:@martin_gorner