compile.ipynb

Kernel: Python 3

Copyright 2019 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

将 XLA 与 tf.function 结合使用

在 TensorFlow.org 查看

在 Google Colab 中运行

在 GitHub 中查看源代码

本教程将训练一个 TensorFlow 模型来对 MNIST 数据集进行分类，我们会使用 XLA 编译训练函数。

首先，加载 TensorFlow 并启用 Eager Execution。

# In TF 2.4 jit_compile is called experimental_compile
!pip install tf-nightly

import tensorflow as tf
tf.compat.v1.enable_eager_execution()

随后，定义一些必要的常量并准备 MNIST 数据集。

# Size of each input image, 28 x 28 pixels
IMAGE_SIZE = 28 * 28
# Number of distinct number labels, [0..9]
NUM_CLASSES = 10
# Number of examples in each training batch (step)
TRAIN_BATCH_SIZE = 100
# Number of training steps to run
TRAIN_STEPS = 1000

# Loads MNIST dataset.
train, test = tf.keras.datasets.mnist.load_data()
train_ds = tf.data.Dataset.from_tensor_slices(train).batch(TRAIN_BATCH_SIZE).repeat()

# Casting from raw data to the required datatypes.
def cast(images, labels):
  images = tf.cast(
      tf.reshape(images, [-1, IMAGE_SIZE]), tf.float32)
  labels = tf.cast(labels, tf.int64)
  return (images, labels)

最后，定义模型和优化器。该模型使用单个密集层。

layer = tf.keras.layers.Dense(NUM_CLASSES)
optimizer = tf.keras.optimizers.Adam()

定义训练函数

在训练函数中，您可以使用上面定义的层来获取预测的标签，然后使用优化器来尽可能减小损失的梯度。为了使用 XLA 编译计算，请将其置于 jit_compile=True 的 tf.function 中。

@tf.function(experimental_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))

训练并测试模型

定义训练函数后，请定义模型。

for images, labels in train_ds:
  if optimizer.iterations > TRAIN_STEPS:
    break
  train_mnist(images, labels)

最后，检查准确率：

images, labels = cast(test[0], test[1])
predicted_labels = layer(images)
correct_prediction = tf.equal(tf.argmax(predicted_labels, 1), labels)
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
print("Prediction accuracy after training: %s" % accuracy)

在后台，XLA 编译器将整个 TF 函数编译为 HLO，后者已启用融合优化。使用自省工具，我们可以查看 HLO 代码（“stage”的其他有趣的可能值是优化后的 HLO 的 optimized_hlo 和 Graphviz 计算图的 optimized_hlo_dot）：

print(train_mnist.experimental_get_compiler_ir(images, labels)(stage='hlo'))