Halaman ini diterjemahkan oleh Cloud Translation API.
Switch to English

Gunakan XLA dengan fungsi

Lihat di TensorFlow.org Jalankan di Google Colab Lihat sumber di GitHub

Tutorial ini melatih model TensorFlow untuk mengklasifikasikan dataset MNIST, di mana fungsi pelatihan dikompilasi menggunakan XLA.

Pertama, muat TensorFlow dan aktifkan eksekusi yang diinginkan.

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

Kemudian tentukan beberapa konstanta yang diperlukan dan siapkan dataset 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

Terakhir, tentukan model dan pengoptimal. Model ini menggunakan lapisan padat tunggal.

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

Tentukan fungsi pelatihan

Dalam fungsi pelatihan, Anda mendapatkan label yang diprediksi menggunakan layer yang ditentukan di atas, dan kemudian meminimalkan gradien kerugian menggunakan optimizer. Untuk mengkompilasi perhitungan menggunakan XLA, letakkan di dalam tf.function dengan 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))
 

Latih dan uji modelnya

Setelah Anda menentukan fungsi pelatihan, tentukan modelnya.

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

Dan, akhirnya, periksa akurasinya:

 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)