使用 TF Profiler 分析 tf.data 效能

總覽

本指南假設您已熟悉 TensorFlow Profilertf.data。本指南旨在提供逐步操作說明和範例,協助使用者診斷並修正輸入管道效能問題。

首先,收集 TensorFlow 作業的設定檔。關於如何操作的說明適用於 CPU/GPUCloud TPU

TensorFlow Trace Viewer

以下詳述的分析工作流程著重於 Profiler 中的追蹤檢視器工具。此工具會顯示時間軸,指出 TensorFlow 程式執行的作業持續時間,並讓您識別哪些作業執行時間最長。如要進一步瞭解追蹤檢視器,請查看 TF Profiler 指南的本節。一般而言,tf.data 事件會顯示在主機 CPU 時間軸上。

分析工作流程

請按照以下工作流程操作。如果您有任何意見可協助我們改善工作流程,請建立 Github 議題,並加上「comp:data」標籤。

1. 您的 tf.data 管道產生資料的速度是否夠快?

首先確定輸入管道是否為 TensorFlow 程式的瓶頸。

若要執行此操作,請在追蹤檢視器中尋找 IteratorGetNext::DoCompute 作業。一般而言,您會預期在步驟開始時看到這些作業。這些切片代表當要求批次元素時,輸入管道產生批次元素所需的時間。如果您使用 keras 或在 tf.function 中疊代資料集,應在 tf_data_iterator_get_next 執行緒中找到這些切片。

請注意,如果您使用分配策略,您可能會看到 IteratorGetNextAsOptional::DoCompute 事件,而不是 IteratorGetNext::DoCompute (截至 TF 2.3)。

image

如果呼叫傳回速度很快 (<= 50 微秒), 表示您的資料在要求時可用。輸入管道不是您的瓶頸;請參閱 Profiler 指南,取得更通用的效能分析提示。

image

如果呼叫傳回速度很慢, tf.data 無法跟上取用者的要求。繼續下一節。

2. 您是否正在預先擷取資料?

輸入管道效能的最佳做法是在 tf.data 管道的結尾插入 tf.data.Dataset.prefetch 轉換。此轉換會將輸入管道的預先處理運算與模型運算的下一個步驟重疊,而且是在訓練模型時獲得最佳輸入管道效能的必要條件。如果您正在預先擷取資料,您應該會在與 IteratorGetNext::DoCompute 作業相同的執行緒上看到 Iterator::Prefetch 切片。

image

如果您在管道結尾沒有 prefetch 則應新增一個。如要進一步瞭解 tf.data 效能建議,請參閱 tf.data 效能指南

如果您已在預先擷取資料,且輸入管道仍是瓶頸,請繼續下一節,進一步分析效能。

3. 您是否達到高 CPU 使用率?

tf.data 藉由盡可能充分利用可用資源來實現高輸送量。一般而言,即使在 GPU 或 TPU 等加速器上執行模型,tf.data 管道也會在 CPU 上執行。您可以使用 sarhtop 等工具,或者如果您在 GCP 上執行,則可在雲端監控主控台中查看您的使用率。

如果您的使用率偏低, 表示您的輸入管道可能未充分利用主機 CPU。您應參閱 tf.data 效能指南以取得最佳做法。如果您已套用最佳做法,且使用率和輸送量仍然偏低,請繼續下方的瓶頸分析

如果您的使用率接近資源限制,為了進一步提升效能,您需要提升輸入管道的效率 (例如,避免不必要的運算) 或卸載運算。

您可以藉由避免在 tf.data 中進行不必要的運算來提升輸入管道的效率。其中一種方法是在運算密集型工作後插入 tf.data.Dataset.cache 轉換 (如果您的資料適合放入記憶體);這會降低運算量,但會增加記憶體使用量。此外,停用 tf.data 中的運算元內平行處理,有可能將效率提升 10% 以上,而且可以透過在您的輸入管道上設定以下選項來完成

dataset = ...
options = tf.data.Options()
options.experimental_threading.max_intra_op_parallelism = 1
dataset = dataset.with_options(options)

4. 瓶頸分析

以下章節將逐步說明如何在追蹤檢視器中讀取 tf.data 事件,以瞭解瓶頸所在以及可能的緩解策略。

瞭解 Profiler 中的 tf.data 事件

Profiler 中的每個 tf.data 事件都有名稱 Iterator::<Dataset>,其中 <Dataset> 是資料集來源或轉換的名稱。每個事件也都有長名稱 Iterator::<Dataset_1>::...::<Dataset_n>,您可以按一下 tf.data 事件來查看長名稱。在長名稱中,<Dataset_n> 符合 (短) 名稱中的 <Dataset>,而長名稱中的其他資料集代表下游轉換。

image

例如,上述螢幕擷取畫面是從以下程式碼產生

dataset = tf.data.Dataset.range(10)
dataset = dataset.map(lambda x: x)
dataset = dataset.repeat(2)
dataset = dataset.batch(5)

在此範例中,Iterator::Map 事件具有長名稱 Iterator::BatchV2::FiniteRepeat::Map。請注意,資料集名稱可能與 Python API 略有不同 (例如,FiniteRepeat 而非 Repeat),但應夠直覺而可剖析。

同步和非同步轉換

對於同步 tf.data 轉換 (例如 BatchMap),您會在相同執行緒上看到來自上游轉換的事件。在上述範例中,由於使用的所有轉換都是同步的,因此所有事件都顯示在相同執行緒上。

對於非同步轉換 (例如 PrefetchParallelMapParallelInterleaveMapAndBatch),來自上游轉換的事件會在不同的執行緒上。在這種情況下,「長名稱」可以協助您識別管道中事件對應的轉換。

image

例如,上述螢幕擷取畫面是從以下程式碼產生

dataset = tf.data.Dataset.range(10)
dataset = dataset.map(lambda x: x)
dataset = dataset.repeat(2)
dataset = dataset.batch(5)
dataset = dataset.prefetch(1)

在此範例中,Iterator::Prefetch 事件位於 tf_data_iterator_get_next 執行緒上。由於 Prefetch 是非同步的,因此其輸入事件 (BatchV2) 會位於不同的執行緒上,而且可以透過搜尋長名稱 Iterator::Prefetch::BatchV2 來找到。在此情況下,它們位於 tf_data_iterator_resource 執行緒上。從其長名稱,您可以推斷 BatchV2Prefetch 的上游。此外,BatchV2 事件的 parent_id 會符合 Prefetch 事件的 ID。

識別瓶頸

一般而言,若要識別輸入管道中的瓶頸,請從最外層的轉換一路到來源巡視輸入管道。從管道中的最後一個轉換開始,遞迴到上游轉換,直到找到速度緩慢的轉換或到達來源資料集 (例如 TFRecord) 為止。在上述範例中,您會從 Prefetch 開始,然後向上游巡視到 BatchV2FiniteRepeatMap,最後到 Range

一般而言,速度緩慢的轉換對應於事件時間很長,但輸入事件時間很短的轉換。以下是一些範例。

請注意,大多數主機輸入管道中的最後一個 (最外層) 轉換是 Iterator::Model 事件。Model 轉換是由 tf.data 執行階段自動導入,用於檢測和自動調整輸入管道效能。

如果您的作業使用分配策略,追蹤檢視器將包含對應於裝置輸入管道的其他事件。裝置管道的最外層轉換 (巢狀於 IteratorGetNextOp::DoComputeIteratorGetNextAsOptionalOp::DoCompute 下方) 將是 Iterator::Prefetch 事件,並具有上游 Iterator::Generator 事件。您可以搜尋 Iterator::Model 事件來找到對應的主機管道。

範例 1

image

上述螢幕擷取畫面是從以下輸入管道產生

dataset = tf.data.TFRecordDataset(filename)
dataset = dataset.map(parse_record)
dataset = dataset.batch(32)
dataset = dataset.repeat()

在螢幕擷取畫面中,觀察到 (1) Iterator::Map 事件時間很長,但 (2) 其輸入事件 (Iterator::FlatMap) 傳回速度很快。這表示循序 Map 轉換是瓶頸。

請注意,在螢幕擷取畫面中,InstantiatedCapturedFunction::Run 事件對應於執行 map 函式所需的時間。

範例 2

image

上述螢幕擷取畫面是從以下輸入管道產生

dataset = tf.data.TFRecordDataset(filename)
dataset = dataset.map(parse_record, num_parallel_calls=2)
dataset = dataset.batch(32)
dataset = dataset.repeat()

此範例與上述範例類似,但使用 ParallelMap 而非 Map。我們在此注意到 (1) Iterator::ParallelMap 事件時間很長,但 (2) 其輸入事件 Iterator::FlatMap (由於 ParallelMap 是非同步的,因此位於不同的執行緒上) 時間很短。這表示 ParallelMap 轉換是瓶頸。

解決瓶頸

來源資料集

如果您已將資料集來源識別為瓶頸 (例如從 TFRecord 檔案讀取),您可以藉由平行化資料擷取來提升效能。若要執行此操作,請確保您的資料已分散到多個檔案中,並搭配 tf.data.Dataset.interleave 使用,並將 num_parallel_calls 參數設定為 tf.data.AUTOTUNE。如果決定性對您的程式而言並不重要,您可以進一步提升效能,方法是在 TF 2.2 版或更新版本中,於 tf.data.Dataset.interleave 上設定 deterministic=False 標記。例如,如果您要從 TFRecord 讀取,您可以執行以下操作

dataset = tf.data.Dataset.from_tensor_slices(filenames)
dataset = dataset.interleave(tf.data.TFRecordDataset,
  num_parallel_calls=tf.data.AUTOTUNE,
  deterministic=False)

請注意,分散式檔案應相當大,以攤銷開啟檔案的額外負荷。如要進一步瞭解平行資料擷取,請參閱 tf.data 效能指南的本節。

轉換資料集

如果您已將中繼 tf.data 轉換識別為瓶頸,您可以藉由平行化轉換或快取運算 (如果您的資料適合放入記憶體且適當) 來解決此問題。有些轉換 (例如 Map) 具有平行對應項目;tf.data 效能指南示範如何平行化這些項目。其他轉換 (例如 FilterUnbatchBatch) 本質上是循序的;您可以藉由導入「外部平行處理」來平行化它們。例如,假設您的輸入管道最初看起來像以下這樣,其中 Batch 作為瓶頸

filenames = tf.data.Dataset.list_files(file_path, shuffle=is_training)
dataset = filenames_to_dataset(filenames)
dataset = dataset.batch(batch_size)

您可以藉由在分散式輸入上執行輸入管道的多個副本並組合結果來導入「外部平行處理」

filenames = tf.data.Dataset.list_files(file_path, shuffle=is_training)

def make_dataset(shard_index):
  filenames = filenames.shard(NUM_SHARDS, shard_index)
  dataset = filenames_to_dataset(filenames)
  Return dataset.batch(batch_size)

indices = tf.data.Dataset.range(NUM_SHARDS)
dataset = indices.interleave(make_dataset,
                             num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.prefetch(tf.data.AUTOTUNE)

其他資源