jit_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.

Cómo usar XLA con tf.function

Ver en TensorFlow.org

Ejecutar en Google Colab

Descargar el bloc de notas

Ver fuente en GitHub

En este tutorial se entrena un modelo de TensorFlow para clasificar el conjunto de datos MNIST, donde la función de entrenamiento se compila con XLA.

Primero, cargue TensorFlow y habilite la ejecución eager.

import tensorflow as tf

Luego, defina algunas constantes necesarias y prepare el conjunto de datos 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)

Finalmente, defina el modelo y el optimizador. El modelo utiliza una única capa densa.

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

Cómo definir la función de entrenamiento

En la función de entrenamiento, se obtienen las etiquetas predichas por medio de la capa definida anteriormente y luego se usa el optimizador para minimizar el gradiente de pérdida. Para compilar el cálculo con XLA, colóquelo dentro de tf.function con jit_compile=True.

@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))

Cómo entrenar y probar el modelo

Una vez que haya definido la función de entrenamiento, defina el modelo.

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

Y, finalmente, verifique la precisión:

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)

En segundo plano, el compilador XLA se encarga de compilar toda la función de TF en HLO, lo que permite optimizaciones de fusión. Por medio de las funciones de introspección, podemos ver el código de HLO (otros valores posibles interesantes para "etapa" son optimized_hlo para HLO después de las optimizaciones y optimized_hlo_dot para un gráfico Graphviz):

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