OpenXLA 是特定領域專用的線性代數編譯器,可在未變更原始碼的情況下,加速 TensorFlow 模型。
引言
執行 TensorFlow 程式時,TensorFlow 執行程式會個別執行所有作業。每項 TensorFlow 作業都有預先編譯的 GPU 核心實作項目,做為執行程式的分派目標。
XLA 提供另一種執行模型的模式:將 TensorFlow 圖形編譯為一系列專為特定模型產生的運算核心。這些核心專屬於模型,因此可以利用模型特定的資訊進行最佳化。我們來看看 XLA 在簡易 TensorFlow 運算中所執行的最佳化作業:
def model_fn(x, y, z):
return tf.reduce_sum(x + y * z)
在沒有 XLA 的情況下執行,圖表會啟動三個核心:一個用於乘法,一個用於加法,另一個用於縮減。但 XLA 可將圖形最佳化,因此只需啟動單一核心即可計算結果。做法是將加法、乘法和減法「融合」為一個 GPU 核心。此外,這個融合作業不會將 y*z
和 x+y*z
產生的中間值寫入記憶體,而是會將這些中間運算的結果直接「串流」給使用者,同時完整保留在 GPU 暫存器中。融合是 XLA 最重要的最佳化功能。記憶體頻寬通常是硬體加速器中最稀少的資源,因此改善效能的最佳方式之一就是移除記憶體運算。
為 TensorFlow 模型啟用 XLA
使用 tf.function(jit_compile=True)
進行明確編譯
明確編譯 API 可讓您精細地控管要選擇哪些函式進行編譯。例如,下列執行 MNIST 訓練的 TensorFlow 函式是使用 XLA 編譯而成:
@tf.function(jit_compile=True)
def train_mnist(images, labels):
images, labels = cast(images, labels)
with tf.GradientTape() as tape:
predicted_labels = layer(images)
loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=predicted_labels, labels=labels
))
layer_variables = layer.trainable_variables
grads = tape.gradient(loss, layer_variables)
optimizer.apply_gradients(zip(grads, layer_variables))
jit_compile
API 具有 must-compile 語意:不是使用 XLA 編譯整個函式,就是擲回 errors.InvalidArgumentError
例外狀況。XLA 目前無法編譯維度無法「推論」的函式 (也就是無法在不執行完整運算的情況下推論所有張量的尺寸)。舉例來說,下列函式無法編譯:
@tf.function
def not_compilable(x):
return tf.unique(x)
不過,形狀在執行過程中可能會有所差異:
@tf.function(jit_compile=True)
def recompiled_on_launch(a, b):
return a + b
recompiled_on_launch(tf.ones([1, 10]), tf.ones([1, 10]))
recompiled_on_launch(tf.ones([1, 100]), tf.ones([1, 100]))
如需更詳細的使用範例,請參閱教學課程 Colab,以及使用 jit_compile=True
的教學影片。
搭配 Keras 使用
針對 Keras 模型,jit_compile=True
可設為 model.compile
的引數:
model.compile(optimizer="adam", jit_compile=True)
採用分散式策略的用量
XLA:GPU 可透過 jit_compile=True
為步驟函式加上註解,以便與 TF 分散式策略 (MirroredStrategy
或 MultiWorkerMirroredStrategy
) 搭配使用:
@tf.function(jit_compile=True)
def step_fn():
t = tf.ones(shape=[100], dtype=tf.float32)
ctx = tf.distribute.get_replica_context()
return ctx.all_reduce(tf.distribute.ReduceOp.SUM, t)
@tf.function
def run_fn():
return strategy.run(step_fn)
自動分群法
如要在 TensorFlow 模型中開始使用 XLA,最簡單的方法就是啟用自動分群功能。這項功能會在可使用 XLA 編譯和執行的 TensorFlow 函式中,自動尋找叢集 (已連結的子圖)。您可以設定 TF_XLA_FLAGS
環境變數,啟用 GPU 上的自動分群法:
$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program
自動分群法目前最適合用於 GPU 工作負載,但您也可以另外使用 --tf_xla_cpu_global_jit
標記,在 CPU 上啟用自動分群法:
$ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program
如需詳細的使用範例,請參閱自動分群法教學課程 Colab。
使用 tfcompile
執行 CPU 的 AOT (預先) 編譯
您也可以使用獨立的 tfcompile
工具,將 TensorFlow 圖形轉換為可執行的程式碼 (僅適用於 x86-64 CPU)。
檢查已編譯的程式
XLA 提供自我檢查設施,可讓您檢查產生的程式。如要傾印產生的程式,請使用環境變數 XLA_FLAGS
:
$ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program
執行傾印後,您可以在 /tmp/generated
中找到下列檔案:
module_XXXX.*_optimizations.txt
產生的 XLA 程式,每個編譯叢集各有一個程式。在提交 XLA 錯誤報告時附上這些內容,非常有幫助!module_XXXX.ptx
產生的 PTX 檔案。
您也可以使用下列方式,傾印以視覺化方式呈現 TensorFlow 圖形中 XLA 叢集嵌入的圖表:
$ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"
可重現的錯誤報告
如果錯誤報告包含產生的 XLA 程式傾印和使用的自動分群嵌入功能,可更輕鬆地重現錯誤報告。如要為執行自動分群功能的 TensorFlow 程式產生這些 API,請啟動:
$ TF_DUMP_GRAPH_PREFIX=/tmp/generated \
TF_XLA_FLAGS="--tf_xla_clustering_debug --tf_xla_auto_jit=2" \
XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=/tmp/generated" \
my/tensorflow/program"
回報錯誤時,請附加 /tmp/generated
目錄的內容 (如上所示)。
請盡可能使用 run_hlo_module
,將錯誤隔離到單一 XLA 程式,並在產生的程式上疊代執行該程式。
延伸閱讀
- OpenXLA 說明文件 OpenXLA 說明文件
- 已知問題:XLA+TF 已知問題清單
- XLA - TensorFlow, Compiled:前往 Google Developers 網誌閱讀
- 查看 GitHub 上的 XLA 原始碼!
XLA 前端
除了 TensorFlow 以外,還可以透過下列方式產生 XLA 程式: