Esta página foi traduzida pela API Cloud Translation.
Switch to English

Use XLA com tf.function

Ver em TensorFlow.org Executar no Google Colab Ver fonte no GitHub

Este tutorial treina um modelo TensorFlow para classificar o conjunto de dados MNIST, em que a função de treinamento é compilada usando XLA.

Primeiro, carregue o TensorFlow e permita uma execução rápida.

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

Em seguida, defina algumas constantes necessárias e prepare o conjunto de dados 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

Por fim, defina o modelo e o otimizador. O modelo usa uma única camada densa.

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

Definir a função de treinamento

Na função de treinamento, você obtém os rótulos previstos usando a camada definida acima e minimiza o gradiente da perda usando o otimizador. Para compilar o cálculo usando XLA, coloque-o dentro de tf.function com 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))
 

Treine e teste o modelo

Depois de definir a função de treinamento, defina o modelo.

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

E, finalmente, verifique a precisão:

 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)