Se usó la API de Cloud Translation para traducir esta página.
Switch to English

Use XLA con tf.function

Ver en TensorFlow.org Ejecutar en Google Colab Ver código fuente en GitHub

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

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

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

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)
 
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11493376/11490434 [==============================] - 1s 0us/step

Finalmente, defina el modelo y el optimizador. El modelo usa una sola capa densa.

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

Definir la función de entrenamiento.

En la función de entrenamiento, obtienes las etiquetas predichas usando la capa definida anteriormente, y luego minimizas el gradiente de la pérdida usando el optimizador. Para compilar el cálculo utilizando XLA, colóquelo dentro de tf.function con experimental_compile=True .

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

Entrena y prueba 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)
 
Prediction accuracy after training: tf.Tensor(0.8829, shape=(), dtype=float32)