Questa pagina è stata tradotta dall'API Cloud Translation.
Switch to English

Usa XLA con la funzione tf

Visualizza su TensorFlow.org Esegui in Google Colab Visualizza sorgente su GitHub

Questo tutorial addestra un modello TensorFlow per classificare il set di dati MNIST, in cui la funzione di addestramento viene compilata utilizzando XLA.

Innanzitutto, carica TensorFlow e abilita l'esecuzione impaziente.

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

Quindi definire alcune costanti necessarie e preparare il set di dati 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

Infine, definisci il modello e l'ottimizzatore. Il modello utilizza un unico strato denso.

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

Definisci la funzione di allenamento

Nella funzione di allenamento, ottieni le etichette previste utilizzando il livello definito sopra, quindi riduci al minimo il gradiente della perdita utilizzando l'ottimizzatore. Per compilare il calcolo usando XLA, posizionalo all'interno di 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))

Addestra e testa il modello

Una volta definita la funzione di addestramento, definire il modello.

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

E, infine, controlla la precisione:

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)