 View on TensorFlow.org
|
 Run in Google Colab
|
 View source on GitHub
|
 Download notebook
|
This tutorial trains a TensorFlow model to classify the CIFAR-10 dataset, and we compile it using XLA.
You will load and normalize the dataset using the TensorFlow Datasets (TFDS) API. First, install/upgrade TensorFlow and TFDS:
pip install -U -q tensorflow tensorflow_datasetsimport tensorflow as tf
import tensorflow_datasets as tfds
# Check that GPU is available: cf. https://colab.research.google.com/notebooks/gpu.ipynb
assert(tf.test.gpu_device_name())
tf.keras.backend.clear_session()
tf.config.optimizer.set_jit(False) # Start with XLA disabled.
def load_data():
result = tfds.load('cifar10', batch_size = -1)
(x_train, y_train) = result['train']['image'],result['train']['label']
(x_test, y_test) = result['test']['image'],result['test']['label']
x_train = x_train.numpy().astype('float32') / 256
x_test = x_test.numpy().astype('float32') / 256
# Convert class vectors to binary class matrices.
y_train = tf.keras.utils.to_categorical(y_train, num_classes=10)
y_test = tf.keras.utils.to_categorical(y_test, num_classes=10)
return ((x_train, y_train), (x_test, y_test))
(x_train, y_train), (x_test, y_test) = load_data()
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR I0000 00:00:1785236877.895634 13197 gpu_device.cc:2020] Created device /device:GPU:0 with 13638 MB memory: -> device: 0, name: Tesla T4, pci bus id: 0000:00:05.0, compute capability: 7.5 I0000 00:00:1785236877.897946 13197 gpu_device.cc:2020] Created device /device:GPU:1 with 13756 MB memory: -> device: 1, name: Tesla T4, pci bus id: 0000:00:06.0, compute capability: 7.5 I0000 00:00:1785236877.900168 13197 gpu_device.cc:2020] Created device /device:GPU:2 with 13756 MB memory: -> device: 2, name: Tesla T4, pci bus id: 0000:00:07.0, compute capability: 7.5 I0000 00:00:1785236877.902403 13197 gpu_device.cc:2020] Created device /device:GPU:3 with 13756 MB memory: -> device: 3, name: Tesla T4, pci bus id: 0000:00:08.0, compute capability: 7.5 I0000 00:00:1785236879.191733 13197 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 13638 MB memory: -> device: 0, name: Tesla T4, pci bus id: 0000:00:05.0, compute capability: 7.5 I0000 00:00:1785236879.193508 13197 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 13756 MB memory: -> device: 1, name: Tesla T4, pci bus id: 0000:00:06.0, compute capability: 7.5 I0000 00:00:1785236879.195421 13197 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:2 with 13756 MB memory: -> device: 2, name: Tesla T4, pci bus id: 0000:00:07.0, compute capability: 7.5 I0000 00:00:1785236879.197177 13197 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:3 with 13756 MB memory: -> device: 3, name: Tesla T4, pci bus id: 0000:00:08.0, compute capability: 7.5
We define the model, adapted from the Keras CIFAR-10 example:
def generate_model():
return tf.keras.models.Sequential([
tf.keras.layers.Conv2D(32, (3, 3), padding='same', input_shape=x_train.shape[1:]),
tf.keras.layers.Activation('relu'),
tf.keras.layers.Conv2D(32, (3, 3)),
tf.keras.layers.Activation('relu'),
tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
tf.keras.layers.Dropout(0.25),
tf.keras.layers.Conv2D(64, (3, 3), padding='same'),
tf.keras.layers.Activation('relu'),
tf.keras.layers.Conv2D(64, (3, 3)),
tf.keras.layers.Activation('relu'),
tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
tf.keras.layers.Dropout(0.25),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(512),
tf.keras.layers.Activation('relu'),
tf.keras.layers.Dropout(0.5),
tf.keras.layers.Dense(10),
tf.keras.layers.Activation('softmax')
])
model = generate_model()
/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/layers/convolutional/base_conv.py:113: UserWarning: Do not pass an `input_shape`/`input_dim` argument to a layer. When using Sequential models, prefer using an `Input(shape)` object as the first layer in the model instead. super().__init__(activity_regularizer=activity_regularizer, **kwargs)
We train the model using the RMSprop optimizer:
def compile_model(model):
opt = tf.keras.optimizers.RMSprop(learning_rate=0.0001)
model.compile(loss='categorical_crossentropy',
optimizer=opt,
metrics=['accuracy'])
return model
model = compile_model(model)
def train_model(model, x_train, y_train, x_test, y_test, epochs=25):
model.fit(x_train, y_train, batch_size=256, epochs=epochs, validation_data=(x_test, y_test), shuffle=True)
def warmup(model, x_train, y_train, x_test, y_test):
# Warm up the JIT, we do not wish to measure the compilation time.
initial_weights = model.get_weights()
train_model(model, x_train, y_train, x_test, y_test, epochs=1)
model.set_weights(initial_weights)
warmup(model, x_train, y_train, x_test, y_test)
%time train_model(model, x_train, y_train, x_test, y_test)
scores = model.evaluate(x_test, y_test, verbose=1)
print('Test loss:', scores[0])
print('Test accuracy:', scores[1])
7/196 ━━━━━━━━━━━━━━━━━━━━ 4s 23ms/step - accuracy: 0.1075 - loss: 2.3060 I0000 00:00:1785236891.788652 13361 device_compiler.h:196] Compiled cluster using XLA! This line is logged at most once for the lifetime of the process. 196/196 ━━━━━━━━━━━━━━━━━━━━ 15s 44ms/step - accuracy: 0.1661 - loss: 2.1990 - val_accuracy: 0.3234 - val_loss: 1.9198 Epoch 1/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 19ms/step - accuracy: 0.1497 - loss: 2.2420 - val_accuracy: 0.3122 - val_loss: 1.9287 Epoch 2/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 3s 17ms/step - accuracy: 0.3043 - loss: 1.9099 - val_accuracy: 0.3839 - val_loss: 1.7249 Epoch 3/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 3s 17ms/step - accuracy: 0.3663 - loss: 1.7516 - val_accuracy: 0.4170 - val_loss: 1.6488 Epoch 4/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 3s 17ms/step - accuracy: 0.3980 - loss: 1.6658 - val_accuracy: 0.4305 - val_loss: 1.5851 Epoch 5/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 3s 17ms/step - accuracy: 0.4161 - loss: 1.6080 - val_accuracy: 0.4597 - val_loss: 1.5111 Epoch 6/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 3s 17ms/step - accuracy: 0.4402 - loss: 1.5487 - val_accuracy: 0.4774 - val_loss: 1.4544 Epoch 7/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 3s 17ms/step - accuracy: 0.4566 - loss: 1.5042 - val_accuracy: 0.4886 - val_loss: 1.4311 Epoch 8/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 3s 17ms/step - accuracy: 0.4676 - loss: 1.4699 - val_accuracy: 0.5031 - val_loss: 1.3858 Epoch 9/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 3s 17ms/step - accuracy: 0.4838 - loss: 1.4330 - val_accuracy: 0.5067 - val_loss: 1.3894 Epoch 10/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 3s 18ms/step - accuracy: 0.4965 - loss: 1.3970 - val_accuracy: 0.5229 - val_loss: 1.3473 Epoch 11/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 3s 18ms/step - accuracy: 0.5036 - loss: 1.3727 - val_accuracy: 0.5211 - val_loss: 1.3540 Epoch 12/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5157 - loss: 1.3450 - val_accuracy: 0.5340 - val_loss: 1.3205 Epoch 13/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5279 - loss: 1.3272 - val_accuracy: 0.5675 - val_loss: 1.2228 Epoch 14/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5398 - loss: 1.2916 - val_accuracy: 0.5665 - val_loss: 1.2318 Epoch 15/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5494 - loss: 1.2714 - val_accuracy: 0.5818 - val_loss: 1.1808 Epoch 16/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5594 - loss: 1.2446 - val_accuracy: 0.5721 - val_loss: 1.1909 Epoch 17/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5695 - loss: 1.2207 - val_accuracy: 0.5982 - val_loss: 1.1465 Epoch 18/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5733 - loss: 1.2041 - val_accuracy: 0.5989 - val_loss: 1.1401 Epoch 19/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5808 - loss: 1.1837 - val_accuracy: 0.6048 - val_loss: 1.1191 Epoch 20/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5900 - loss: 1.1603 - val_accuracy: 0.6012 - val_loss: 1.1546 Epoch 21/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5969 - loss: 1.1361 - val_accuracy: 0.6182 - val_loss: 1.0781 Epoch 22/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.6026 - loss: 1.1270 - val_accuracy: 0.6277 - val_loss: 1.0636 Epoch 23/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.6052 - loss: 1.1112 - val_accuracy: 0.6320 - val_loss: 1.0559 Epoch 24/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.6132 - loss: 1.0999 - val_accuracy: 0.6393 - val_loss: 1.0368 Epoch 25/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.6218 - loss: 1.0821 - val_accuracy: 0.6350 - val_loss: 1.0506 CPU times: user 40.9 s, sys: 12.8 s, total: 53.7 s Wall time: 1min 30s 313/313 ━━━━━━━━━━━━━━━━━━━━ 1s 2ms/step - accuracy: 0.6392 - loss: 1.0359 Test loss: 1.050584316253662 Test accuracy: 0.6349999904632568
Now let's train the model again, using the XLA compiler. To enable the compiler in the middle of the application, we need to reset the Keras session.
# We need to clear the session to enable JIT in the middle of the program.
tf.keras.backend.clear_session()
tf.config.optimizer.set_jit(True) # Enable XLA.
model = compile_model(generate_model())
(x_train, y_train), (x_test, y_test) = load_data()
warmup(model, x_train, y_train, x_test, y_test)
%time train_model(model, x_train, y_train, x_test, y_test)
196/196 ━━━━━━━━━━━━━━━━━━━━ 11s 39ms/step - accuracy: 0.1634 - loss: 2.2170 - val_accuracy: 0.3122 - val_loss: 1.9344 Epoch 1/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 20ms/step - accuracy: 0.1467 - loss: 2.2488 - val_accuracy: 0.3094 - val_loss: 1.9355 Epoch 2/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.3058 - loss: 1.9148 - val_accuracy: 0.3950 - val_loss: 1.7146 Epoch 3/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.3714 - loss: 1.7379 - val_accuracy: 0.4241 - val_loss: 1.6110 Epoch 4/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.3988 - loss: 1.6459 - val_accuracy: 0.4263 - val_loss: 1.6281 Epoch 5/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.4250 - loss: 1.5741 - val_accuracy: 0.4782 - val_loss: 1.4504 Epoch 6/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.4541 - loss: 1.5092 - val_accuracy: 0.4810 - val_loss: 1.4258 Epoch 7/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 19ms/step - accuracy: 0.4606 - loss: 1.4789 - val_accuracy: 0.4832 - val_loss: 1.4156 Epoch 8/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 19ms/step - accuracy: 0.4851 - loss: 1.4300 - val_accuracy: 0.5145 - val_loss: 1.3592 Epoch 9/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.4964 - loss: 1.4010 - val_accuracy: 0.5195 - val_loss: 1.3634 Epoch 10/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5080 - loss: 1.3730 - val_accuracy: 0.5225 - val_loss: 1.3280 Epoch 11/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5221 - loss: 1.3495 - val_accuracy: 0.5398 - val_loss: 1.2799 Epoch 12/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5247 - loss: 1.3223 - val_accuracy: 0.5181 - val_loss: 1.3790 Epoch 13/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5343 - loss: 1.2948 - val_accuracy: 0.5647 - val_loss: 1.2213 Epoch 14/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5419 - loss: 1.2808 - val_accuracy: 0.5713 - val_loss: 1.2065 Epoch 15/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5535 - loss: 1.2564 - val_accuracy: 0.5852 - val_loss: 1.1739 Epoch 16/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5624 - loss: 1.2390 - val_accuracy: 0.5947 - val_loss: 1.1564 Epoch 17/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5732 - loss: 1.2065 - val_accuracy: 0.5768 - val_loss: 1.1842 Epoch 18/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5767 - loss: 1.1947 - val_accuracy: 0.5587 - val_loss: 1.2439 Epoch 19/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5842 - loss: 1.1764 - val_accuracy: 0.6219 - val_loss: 1.0885 Epoch 20/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.5920 - loss: 1.1529 - val_accuracy: 0.6102 - val_loss: 1.1136 Epoch 21/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.6036 - loss: 1.1233 - val_accuracy: 0.6259 - val_loss: 1.0736 Epoch 22/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.6101 - loss: 1.1089 - val_accuracy: 0.6301 - val_loss: 1.0471 Epoch 23/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.6171 - loss: 1.0926 - val_accuracy: 0.6355 - val_loss: 1.0415 Epoch 24/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.6203 - loss: 1.0825 - val_accuracy: 0.6309 - val_loss: 1.0456 Epoch 25/25 196/196 ━━━━━━━━━━━━━━━━━━━━ 4s 18ms/step - accuracy: 0.6284 - loss: 1.0659 - val_accuracy: 0.6529 - val_loss: 0.9915 CPU times: user 41.3 s, sys: 12.4 s, total: 53.7 s Wall time: 1min 32s
On a machine with a Titan V GPU and an Intel Xeon E5-2690 CPU the speed up is ~1.17x.